Update landmarkdiff/data.py to v0.3.2

landmarkdiff/data.py

@@ -38,7 +38,6 @@ logger = logging.getLogger(__name__)
 # Core dataset
 # ---------------------------------------------------------------------------
 
-
 class SurgicalPairDataset(Dataset):
     """Dataset for loading surgical before/after training pairs.
 
@@ -163,7 +162,9 @@ class SurgicalPairDataset(Dataset):
         img = cv2.imread(str(path))
         if img is None:
             logger.warning("Failed to load %s, using blank", path)
-            return np.zeros((self.resolution, self.resolution, 3), dtype=np.uint8)
+            return np.zeros(
+                (self.resolution, self.resolution, 3), dtype=np.uint8
+            )
         if img.shape[:2] != (self.resolution, self.resolution):
             img = cv2.resize(img, (self.resolution, self.resolution))
         return img
@@ -172,10 +173,14 @@ class SurgicalPairDataset(Dataset):
         """Load a mask as float32 [0,1], resized to resolution."""
         path = self.data_dir / filename
         if not path.exists():
-            return np.ones((self.resolution, self.resolution), dtype=np.float32)
+            return np.ones(
+                (self.resolution, self.resolution), dtype=np.float32
+            )
         mask = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE)
         if mask is None:
-            return np.ones((self.resolution, self.resolution), dtype=np.float32)
+            return np.ones(
+                (self.resolution, self.resolution), dtype=np.float32
+            )
         mask = cv2.resize(mask, (self.resolution, self.resolution))
         return mask.astype(np.float32) / 255.0
 
@@ -184,7 +189,6 @@ class SurgicalPairDataset(Dataset):
 # Evaluation dataset (input + ground truth)
 # ---------------------------------------------------------------------------
 
-
 class EvalPairDataset(Dataset):
     """Dataset for evaluation: loads input/target pairs with procedure labels.
 
@@ -231,7 +235,9 @@ class EvalPairDataset(Dataset):
         path = self.data_dir / filename
         img = cv2.imread(str(path))
         if img is None:
-            return np.zeros((self.resolution, self.resolution, 3), dtype=np.uint8)
+            return np.zeros(
+                (self.resolution, self.resolution, 3), dtype=np.uint8
+            )
         if img.shape[:2] != (self.resolution, self.resolution):
             img = cv2.resize(img, (self.resolution, self.resolution))
         return img
@@ -241,7 +247,6 @@ class EvalPairDataset(Dataset):
 # Conversion utilities
 # ---------------------------------------------------------------------------
 
-
 def bgr_to_tensor(bgr: np.ndarray) -> torch.Tensor:
     """Convert BGR uint8 image to RGB [0,1] tensor (C, H, W)."""
     rgb = bgr[:, :, ::-1].astype(np.float32) / 255.0
@@ -266,7 +271,6 @@ def mask_to_tensor(mask: np.ndarray) -> torch.Tensor:
 # Samplers
 # ---------------------------------------------------------------------------
 
-
 def create_procedure_sampler(
     dataset: SurgicalPairDataset,
     balance_procedures: bool = True,
@@ -305,7 +309,6 @@ def create_procedure_sampler(
 # DataLoader factory
 # ---------------------------------------------------------------------------
 
-
 def create_dataloader(
     dataset: Dataset,
     batch_size: int = 4,
@@ -350,7 +353,6 @@ def create_dataloader(
 # Multi-directory dataset
 # ---------------------------------------------------------------------------
 
-
 class CombinedDataset(Dataset):
     """Combine multiple SurgicalPairDatasets into one.
 
@@ -372,6 +374,10 @@ class CombinedDataset(Dataset):
         return self._cumulative_sizes[-1] if self._cumulative_sizes else 0
 
     def __getitem__(self, idx: int) -> dict:
+        if idx < 0 or idx >= len(self):
+            raise IndexError(
+                f"CombinedDataset index {idx} out of range [0, {len(self)})"
+            )
         dataset_idx = 0
         for i, size in enumerate(self._cumulative_sizes):
             if idx < size:
@@ -382,6 +388,10 @@ class CombinedDataset(Dataset):
         return self.datasets[dataset_idx][idx]
 
     def get_procedure(self, idx: int) -> str:
+        if idx < 0 or idx >= len(self):
+            raise IndexError(
+                f"CombinedDataset index {idx} out of range [0, {len(self)})"
+            )
         dataset_idx = 0
         for i, size in enumerate(self._cumulative_sizes):
             if idx < size: