Add dataset classes (pseudo-label + synthetic)

dataset.py

@@ -0,0 +1,292 @@
+"""
+Dataset and augmentation for pseudo-label distillation training.
+
+Loads image + pseudo-label pairs (mask, flows, distance transform)
+generated by Stage 1 (generate_pseudolabels.py).
+
+Augmentation strategy (from Cellpose training + PicoSAM2):
+  - Random horizontal/vertical flips
+  - Random rotation (0, 90, 180, 270)
+  - Random crop (if images are large)
+  - Intensity jitter (brightness, contrast)
+  - Gaussian noise
+  - Elastic deformation (circles are robust to small deformations)
+"""
+
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from scipy import ndimage
+from skimage import io as skio
+
+
+class BubblePseudoLabelDataset(Dataset):
+    """
+    Dataset for distillation training from Cellpose pseudo-labels.
+
+    Expects a directory with files:
+      {image_stem}_mask.npy   — instance mask (H, W), int32
+      {image_stem}_flows.npy  — flow fields (3, H, W), float32 [dY, dX, cell_prob]
+      {image_stem}_dist.npy   — distance transform (H, W), float32
+
+    And the original images in a separate directory.
+
+    Args:
+        image_dir: Directory with original images
+        label_dir: Directory with pseudo-labels from Stage 1
+        crop_size: Random crop size (H, W). None for full images.
+        augment: Whether to apply data augmentation
+        normalize: Whether to normalize images to [0, 1]
+    """
+
+    def __init__(
+        self,
+        image_dir,
+        label_dir,
+        crop_size=None,
+        augment=True,
+        normalize=True,
+    ):
+        self.image_dir = Path(image_dir)
+        self.label_dir = Path(label_dir)
+        self.crop_size = crop_size
+        self.augment = augment
+        self.normalize = normalize
+
+        # Find matching image-label pairs
+        extensions = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"}
+        image_files = {f.stem: f for f in self.image_dir.iterdir() if f.suffix.lower() in extensions}
+        label_stems = {f.stem.replace("_mask", "") for f in self.label_dir.glob("*_mask.npy")}
+
+        self.samples = []
+        for stem in sorted(label_stems):
+            if stem in image_files:
+                self.samples.append({
+                    "image_path": image_files[stem],
+                    "mask_path": self.label_dir / f"{stem}_mask.npy",
+                    "flows_path": self.label_dir / f"{stem}_flows.npy",
+                    "dist_path": self.label_dir / f"{stem}_dist.npy",
+                })
+
+        print(f"Found {len(self.samples)} image-label pairs")
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        sample = self.samples[idx]
+
+        # Load image
+        img = skio.imread(str(sample["image_path"]))
+        if img.ndim == 3 and img.shape[2] >= 3:
+            img = np.mean(img[:, :, :3], axis=2)  # to grayscale
+        img = img.astype(np.float32)
+
+        # Load pseudo-labels
+        mask = np.load(sample["mask_path"]).astype(np.float32)
+        flows = np.load(sample["flows_path"]).astype(np.float32)  # (3, H, W)
+        dist = np.load(sample["dist_path"]).astype(np.float32)
+
+        # Normalize image
+        if self.normalize:
+            img = self._normalize_image(img)
+
+        # Normalize targets
+        flow_dY = flows[0]
+        flow_dX = flows[1]
+        cell_prob = flows[2]
+
+        # Normalize cell_prob to [0, 1] if not already
+        if cell_prob.max() > 1.0:
+            cell_prob = (cell_prob - cell_prob.min()) / (cell_prob.max() - cell_prob.min() + 1e-8)
+
+        # Binary mask from instance mask (for cell_prob target)
+        binary_mask = (mask > 0).astype(np.float32)
+
+        # Normalize distance transform: divide by max to get [0, 1] range
+        dist_max = dist.max()
+        if dist_max > 0:
+            dist_norm = dist / dist_max
+        else:
+            dist_norm = dist
+
+        # Stack targets: (4, H, W) = [dY, dX, cell_prob_binary, dist_norm]
+        target = np.stack([flow_dY, flow_dX, binary_mask, dist_norm], axis=0)
+
+        # Augmentation
+        if self.augment:
+            img, target = self._augment(img, target)
+
+        # Random crop
+        if self.crop_size is not None:
+            img, target = self._random_crop(img, target, self.crop_size)
+
+        # To tensors
+        img_tensor = torch.from_numpy(img[np.newaxis]).float()  # (1, H, W)
+        target_tensor = torch.from_numpy(target).float()  # (4, H, W)
+
+        return img_tensor, target_tensor, dist_max  # dist_max for denormalization
+
+    def _normalize_image(self, img):
+        """Percentile normalization (robust to outliers)."""
+        p1, p99 = np.percentile(img, [1, 99])
+        if p99 - p1 > 0:
+            img = (img - p1) / (p99 - p1)
+        else:
+            img = img / (img.max() + 1e-8)
+        return np.clip(img, 0, 1)
+
+    def _augment(self, img, target):
+        """Apply geometric and intensity augmentations.
+
+        Geometric augmentations are applied consistently to image and targets.
+        Intensity augmentations are applied only to the image.
+        """
+        # Random horizontal flip
+        if np.random.random() < 0.5:
+            img = np.flip(img, axis=1).copy()
+            target = np.flip(target, axis=2).copy()
+            target[1] = -target[1]  # flip dX direction
+
+        # Random vertical flip
+        if np.random.random() < 0.5:
+            img = np.flip(img, axis=0).copy()
+            target = np.flip(target, axis=1).copy()
+            target[0] = -target[0]  # flip dY direction
+
+        # Random 90-degree rotations
+        k = np.random.randint(4)
+        if k > 0:
+            img = np.rot90(img, k).copy()
+            target = np.rot90(target, k, axes=(1, 2)).copy()
+            # Adjust flow directions for rotation
+            if k == 1:  # 90° CCW
+                target[0], target[1] = target[1].copy(), -target[0].copy()
+            elif k == 2:  # 180°
+                target[0] = -target[0]
+                target[1] = -target[1]
+            elif k == 3:  # 270° CCW = 90° CW
+                target[0], target[1] = -target[1].copy(), target[0].copy()
+
+        # Intensity augmentations (image only)
+        if np.random.random() < 0.5:
+            delta = np.random.uniform(-0.1, 0.1)
+            img = np.clip(img + delta, 0, 1)
+
+        if np.random.random() < 0.5:
+            factor = np.random.uniform(0.8, 1.2)
+            mean = img.mean()
+            img = np.clip((img - mean) * factor + mean, 0, 1)
+
+        if np.random.random() < 0.3:
+            noise = np.random.normal(0, 0.02, img.shape).astype(np.float32)
+            img = np.clip(img + noise, 0, 1)
+
+        return img, target
+
+    def _random_crop(self, img, target, crop_size):
+        """Random crop of image and target."""
+        h, w = img.shape[-2:] if img.ndim >= 2 else img.shape
+        ch, cw = crop_size
+
+        if h <= ch or w <= cw:
+            pad_h = max(ch - h, 0)
+            pad_w = max(cw - w, 0)
+            img = np.pad(img, ((0, pad_h), (0, pad_w)), mode="reflect")
+            target = np.pad(target, ((0, 0), (0, pad_h), (0, pad_w)), mode="reflect")
+            h, w = img.shape[-2:] if img.ndim >= 2 else img.shape
+
+        y = np.random.randint(0, h - ch + 1)
+        x = np.random.randint(0, w - cw + 1)
+
+        img = img[y : y + ch, x : x + cw]
+        target = target[:, y : y + ch, x : x + cw]
+
+        return img, target
+
+
+class SyntheticBubbleDataset(Dataset):
+    """
+    Synthetic dataset for testing/debugging the training pipeline.
+    Generates random circle images with corresponding labels.
+
+    Useful for verifying the model trains correctly before using real data.
+    """
+
+    def __init__(self, n_samples=100, image_size=256, n_bubbles_range=(5, 30),
+                 radius_range=(5, 25), noise_std=0.05):
+        self.n_samples = n_samples
+        self.image_size = image_size
+        self.n_bubbles_range = n_bubbles_range
+        self.radius_range = radius_range
+        self.noise_std = noise_std
+
+    def __len__(self):
+        return self.n_samples
+
+    def __getitem__(self, idx):
+        H = W = self.image_size
+        img = np.random.uniform(0.1, 0.3, (H, W)).astype(np.float32)
+
+        mask = np.zeros((H, W), dtype=np.int32)
+        n_bubbles = np.random.randint(*self.n_bubbles_range)
+
+        yy, xx = np.mgrid[:H, :W]
+
+        for i in range(1, n_bubbles + 1):
+            cx = np.random.randint(20, W - 20)
+            cy = np.random.randint(20, H - 20)
+            r = np.random.randint(*self.radius_range)
+
+            dist = np.sqrt((yy - cy) ** 2 + (xx - cx) ** 2)
+            circle = dist <= r
+
+            overlap = (mask > 0) & circle
+            if overlap.sum() > 0.3 * circle.sum():
+                continue
+
+            mask[circle] = i
+
+            ring = (dist >= r - 2) & (dist <= r + 1)
+            img[circle] += np.random.uniform(0.3, 0.5)
+            img[ring] += np.random.uniform(0.1, 0.2)
+
+        img += np.random.normal(0, self.noise_std, img.shape).astype(np.float32)
+        img = np.clip(img, 0, 1)
+
+        # Compute targets
+        flow_dY = np.zeros((H, W), dtype=np.float32)
+        flow_dX = np.zeros((H, W), dtype=np.float32)
+        for label_id in range(1, mask.max() + 1):
+            instance = (mask == label_id)
+            if instance.sum() == 0:
+                continue
+            cy_m, cx_m = ndimage.center_of_mass(instance)
+            flow_dY[instance] = cy_m - yy[instance]
+            flow_dX[instance] = cx_m - xx[instance]
+
+        flow_max = max(np.abs(flow_dY).max(), np.abs(flow_dX).max(), 1e-8)
+        flow_dY /= flow_max
+        flow_dX /= flow_max
+
+        binary_mask = (mask > 0).astype(np.float32)
+
+        dist_transform = np.zeros((H, W), dtype=np.float32)
+        for label_id in range(1, mask.max() + 1):
+            instance = (mask == label_id)
+            if instance.sum() == 0:
+                continue
+            d = ndimage.distance_transform_edt(instance)
+            dist_transform[instance] = d[instance]
+        dist_max = dist_transform.max() if dist_transform.max() > 0 else 1.0
+        dist_norm = dist_transform / dist_max
+
+        target = np.stack([flow_dY, flow_dX, binary_mask, dist_norm], axis=0)
+
+        img_tensor = torch.from_numpy(img[np.newaxis]).float()
+        target_tensor = torch.from_numpy(target).float()
+
+        return img_tensor, target_tensor, dist_max