Add all files

Revert too extreme code clean up
biomedia-mira · May 27, 2024 · 5bd5fc9 · 5bd5fc9
1 parent 46b355a
commit 5bd5fc9
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/data_handling/xray.py b/data_handling/xray.py
@@ -276,9 +276,14 @@ def __getitem__(self, idx: int) -> Dict:
             ).detach()
 
         if self.use_counterfactuals:
-            cfx = self.load_counterfactual_image(idx)
-            if not self.counterfactual_contrastive_pairs:
-                img = cfx.clone()
+            if self.counterfactual_contrastive_pairs:
+                cfx = self.load_counterfactual_image(idx)
+            else:
+                if torch.rand(1).item() > 0.5:
+                    cfx = self.load_counterfactual_image(idx)
+                    img = cfx.clone()
+                else:
+                    cfx = img.clone()
             img = self.transform(img)
             cfx = self.transform(cfx)
             img = torch.stack([img, cfx], dim=0).float()