ianpan
/

total-classifier

Image Classification

total_classifier

feature-extraction

Model card Files Files and versions

ianpan commited on Feb 25

Commit

1eebdfb

·

verified ·

1 Parent(s): 1d7a3e1

Upload model

Files changed (1) hide show

modeling.py +2 -2

modeling.py CHANGED Viewed

@@ -510,7 +510,7 @@ class TotalClassifierModel(PreTrainedModel):
         if 1 in crop_dims:
             # swap plane
-            x = x0.transpose(1, 0, 2)
             smin1, smax1 = self.crop_single_plane(
                 x, device, organ, threshold, buffer, speed_up
             )
@@ -519,7 +519,7 @@ class TotalClassifierModel(PreTrainedModel):
         if 2 in crop_dims:
             # swap plane
-            x = x0.transpose(2, 1, 0)
             smin2, smax2 = self.crop_single_plane(
                 x, device, organ, threshold, buffer, speed_up
             )

         if 1 in crop_dims:
             # swap plane
+            x = x0.swapaxes(1, 0)
             smin1, smax1 = self.crop_single_plane(
                 x, device, organ, threshold, buffer, speed_up
             )
         if 2 in crop_dims:
             # swap plane
+            x = x0.transpose(2, 0)
             smin2, smax2 = self.crop_single_plane(
                 x, device, organ, threshold, buffer, speed_up
             )