remove reference to TF.get_dimensions for compatibility (#19)

teamtomo · Feb 22, 2023 · 9dfcbe0 · 9dfcbe0
1 parent e045b28
commit 9dfcbe0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/src/fidder/predict/predict.py b/src/fidder/predict/predict.py
@@ -2,7 +2,6 @@
 from typing import Optional, Tuple
 
 import torch
-import torchvision.transforms.functional as TF
 from einops import rearrange
 from pytorch_lightning import Trainer
 
@@ -38,7 +37,7 @@ def predict_fiducial_mask(
     elif image.ndim == 3:
         image = rearrange(image, "b h w -> b 1 h w")
     image = torch.as_tensor(image, dtype=torch.float)
-    _, h, w = TF.get_dimensions(image)
+    h, w = image.shape[-2:]
     downscale_factor = calculate_resampling_factor(
         source=pixel_spacing, target=TRAINING_PIXEL_SIZE
     )

diff --git a/src/fidder/utils.py b/src/fidder/utils.py
@@ -35,7 +35,7 @@ def rescale_2d_bicubic(
         `(b, c, h, w)` array of rescaled image(s).
     """
     if factor is not None:
-        _, h, w = TF.get_dimensions(image)
+        h, w = image.shape[-2:]
         size = int(factor * min(h, w))
     rescaled_image = TF.resize(
         image, size=size, interpolation=TF.InterpolationMode.BICUBIC
@@ -63,7 +63,7 @@ def rescale_2d_nearest(
         `(b, c, h, w)` array of rescaled image(s).
     """
     if factor is not None:
-        _, h, w = TF.get_dimensions(image)
+        h, w = image.shape[-2:]
         size = int(factor * min(h, w))
     rescaled_image = TF.resize(
         image, size=size, interpolation=TF.InterpolationMode.NEAREST