Upload datasets.py

datasets.py

@@ -316,7 +316,7 @@ def preproc_mimic(batch):
         elif k in ["race"]:
             batch[k] = F.one_hot(batch[k], num_classes=3).squeeze().float()
         elif k in ["finding"]:
-            batch[k] = batch[k].unsqueeze(-1).float()
+            batch[k] = F.one_hot(batch[k], num_classes=3).squeeze().float()
         else:
             batch[k] = batch[k].float().unsqueeze(-1)
     return batch
@@ -336,8 +336,8 @@ class MIMICDataset(Dataset):
         self.transform = transform
         self.disease_labels = [
             "No Finding",
-            "Other",
-            "Pleural Effusion",
+            "Consolidation",
+            "Lung Opacity",
             # "Lung Opacity",
         ]
         self.samples = {
@@ -362,7 +362,16 @@ class MIMICDataset(Dataset):
             # self.samples["pleural_effusion"].append(pleural_effusion)
 
             disease = self.data.loc[idx, "disease"]
-            finding = 0 if disease == "No Finding" else 1
+            # finding = 0 if disease == "No Finding" else 1
+            if disease == "No Finding":
+                finding = 0
+            elif disease == "Consolidation":
+                finding = 1 
+            elif disease == "Lung Opacity":
+                finding = 2
+            else:
+                finding = None
+                    
 
             self.samples["x"].append(img_path)
             self.samples["finding"].append(finding)