Upload model

modeling.py

@@ -162,17 +162,24 @@ class TotalClassifierModel(PreTrainedModel):
             return batch_list
 
         if return_as_list:
-            # list of lists
+            # returns list of list of lists of strings
+            # innermost list - list of strings for each organ present based on threshold
+            # inner list - list of above for each slice
+            # outer list - list of above for each batch element (studies)
             batch_list = []
             for i in range(probas.shape[0]):
                 probas_i = probas[i]
-                batch_list.append(
-                    [
-                        self.index2label[each_class]
-                        for each_class in range(probas_i.shape[1])
-                        if probas_i[:, each_class] >= threshold
-                    ]
-                )
+                for each_slice in range(probas_i.shape[0]):
+                    list_for_batch = []
+                    for each_class in range(probas_i.shape[1]):
+                        list_for_batch.append(
+                            [
+                                self.index2label[each_class]
+                                for each_class in range(probas_i.shape[1])
+                                if probas_i[each_slice, each_class] >= threshold
+                            ]
+                        )
+                    batch_list.append(list_for_batch)
             return batch_list
 
         return probas
@@ -409,6 +416,7 @@ class TotalClassifierModel(PreTrainedModel):
         x: np.ndarray,
         mode: str = "2d",
         torchify: bool = True,
+        add_batch_dim: bool = False,
         device: str | torch.device | None = None,
     ) -> np.ndarray:
         if device is not None:
@@ -427,6 +435,11 @@ class TotalClassifierModel(PreTrainedModel):
                 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")
+        if add_batch_dim:
+            if torchify:
+                x = x.unsqueeze(0)
+            else:
+                x = x[np.newaxis]
         if device is not None:
             x = x.to(device)
         return x