ianpan
/

total-classifier

Image Classification

total_classifier

feature-extraction

Model card Files Files and versions

ianpan commited on Mar 1, 2025

Commit

07c8d8d

·

verified ·

1 Parent(s): 9e9710a

Upload model

Files changed (1) hide show

modeling.py +2 -2

modeling.py CHANGED Viewed

@@ -418,9 +418,9 @@ class TotalClassifierModel(PreTrainedModel):
                 x = x[:, :, :, np.newaxis]
         if torchify:
             if x.ndim == 3:
-                x = torch.from_numpy(x).float().permute(2, 0, 1)
             elif x.ndim == 4:
-                x = torch.from_numpy(x).float().permute(3, 0, 1, 2)
         return x
     def crop_single_plane(

                 x = x[:, :, :, np.newaxis]
         if torchify:
             if x.ndim == 3:
+                x = rearrange(torch.from_numpy(x).float(), "h w c -> c h w")
             elif x.ndim == 4:
+                x = rearrange(torch.from_numpy(x).float(), "n h w c -> n c h w")
         return x
     def crop_single_plane(