Upload src/preprocessing.py with huggingface_hub

src/preprocessing.py

@@ -0,0 +1,284 @@
+"""
+Data loading, annotation parsing, and preprocessing for immunogold TEM images.
+
+The model receives raw images — the CEM500K backbone was pretrained on raw EM.
+Top-hat preprocessing is only used by LodeStar (Stage 1).
+"""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import tifffile
+
+
+# ---------------------------------------------------------------------------
+# Data registry: robust discovery of images, masks, and annotations
+# ---------------------------------------------------------------------------
+
+@dataclass
+class SynapseRecord:
+    """Metadata for one synapse sample."""
+    synapse_id: str
+    image_path: Path
+    mask_path: Optional[Path]
+    csv_6nm_paths: List[Path] = field(default_factory=list)
+    csv_12nm_paths: List[Path] = field(default_factory=list)
+    has_6nm: bool = False
+    has_12nm: bool = False
+
+
+def discover_synapse_data(root: str, synapse_ids: List[str]) -> List[SynapseRecord]:
+    """
+    Discover all TIF images, masks, and CSV annotations for each synapse.
+
+    Handles naming inconsistencies:
+    - S22: main image is S22_0003.tif, two Results folders
+    - S25: 12nm CSV has no space ("Results12nm")
+    - CSV patterns: "Results 6nm XY" vs "Results XY in microns 6nm"
+    """
+    root = Path(root)
+    analyzed = root / "analyzed synapses"
+    records = []
+
+    for sid in synapse_ids:
+        folder = analyzed / sid
+        if not folder.exists():
+            raise FileNotFoundError(f"Synapse folder not found: {folder}")
+
+        # --- Find main image (TIF without 'mask' or 'color' in name) ---
+        all_tifs = list(folder.glob("*.tif"))
+        main_tifs = [
+            t for t in all_tifs
+            if "mask" not in t.stem.lower() and "color" not in t.stem.lower()
+        ]
+        if not main_tifs:
+            raise FileNotFoundError(f"No main image found in {folder}")
+        # Prefer the largest file (main EM image) if multiple found
+        image_path = max(main_tifs, key=lambda t: t.stat().st_size)
+
+        # --- Find mask ---
+        mask_tifs = [t for t in all_tifs if "mask" in t.stem.lower()]
+        mask_path = None
+        if mask_tifs:
+            # Prefer plain "mask.tif" over "mask 1.tif" / "mask 2.tif"
+            plain = [t for t in mask_tifs if t.stem.lower().endswith("mask")]
+            mask_path = plain[0] if plain else mask_tifs[0]
+
+        # --- Find CSVs across all Results* subdirectories ---
+        results_dirs = sorted(folder.glob("Results*"))
+        # Also check direct subdirs like "Results 1", "Results 2"
+        csv_6nm_paths = []
+        csv_12nm_paths = []
+
+        for rdir in results_dirs:
+            if rdir.is_dir():
+                for csv_file in rdir.glob("*.csv"):
+                    name_lower = csv_file.name.lower()
+                    if "6nm" in name_lower:
+                        csv_6nm_paths.append(csv_file)
+                    elif "12nm" in name_lower:
+                        csv_12nm_paths.append(csv_file)
+
+        record = SynapseRecord(
+            synapse_id=sid,
+            image_path=image_path,
+            mask_path=mask_path,
+            csv_6nm_paths=csv_6nm_paths,
+            csv_12nm_paths=csv_12nm_paths,
+            has_6nm=len(csv_6nm_paths) > 0,
+            has_12nm=len(csv_12nm_paths) > 0,
+        )
+        records.append(record)
+
+    return records
+
+
+# ---------------------------------------------------------------------------
+# Image I/O
+# ---------------------------------------------------------------------------
+
+def load_image(path: Path) -> np.ndarray:
+    """
+    Load a TIF image as grayscale uint8.
+
+    Handles:
+    - RGB images (take first channel)
+    - Palette-mode images
+    - Already-grayscale images
+    """
+    img = tifffile.imread(str(path))
+    if img.ndim == 3:
+        # RGB or multi-channel — take first channel (all channels identical in these images)
+        img = img[:, :, 0] if img.shape[2] <= 4 else img[0]
+    return img.astype(np.uint8)
+
+
+def load_mask(path: Path) -> np.ndarray:
+    """
+    Load mask TIF as binary array.
+
+    Mask is RGB where tissue regions have values < 250 in at least one channel.
+    Returns boolean array: True = tissue/structural region.
+    """
+    mask_rgb = tifffile.imread(str(path))
+    if mask_rgb.ndim == 2:
+        return mask_rgb < 250
+    # RGB mask: tissue where any channel is not white
+    return np.any(mask_rgb < 250, axis=-1)
+
+
+# ---------------------------------------------------------------------------
+# Annotation loading and coordinate conversion
+# ---------------------------------------------------------------------------
+
+def load_annotations_csv(csv_path: Path) -> pd.DataFrame:
+    """
+    Load annotation CSV with columns [index, X, Y].
+
+    CSV headers have leading space: " ,X,Y".
+    Coordinates are normalized [0, 1] despite 'microns' in filename.
+    """
+    df = pd.read_csv(csv_path)
+    # Normalize column names (strip whitespace)
+    df.columns = [c.strip() for c in df.columns]
+    # Rename unnamed index column
+    if "" in df.columns:
+        df = df.rename(columns={"": "idx"})
+    return df[["X", "Y"]]
+
+
+# Micron-to-pixel scale factor: consistent across all synapses (verified
+# against researcher's color overlay TIFs). The CSV columns labeled "XY in
+# microns" really ARE microns — multiply by this constant to get pixels.
+MICRONS_TO_PIXELS = 1790.0
+
+
+def load_all_annotations(
+    record: SynapseRecord, image_shape: Tuple[int, int]
+) -> Dict[str, np.ndarray]:
+    """
+    Load and convert annotations for one synapse to pixel coordinates.
+
+    CSV coordinates are in microns (despite filename suggesting normalization).
+    Multiply by MICRONS_TO_PIXELS (1790 px/micron) to convert.
+
+    Args:
+        record: SynapseRecord with CSV paths.
+        image_shape: (height, width) of the corresponding image.
+
+    Returns:
+        Dictionary with keys '6nm' and '12nm', each containing
+        an Nx2 array of (x, y) pixel coordinates.
+    """
+    h, w = image_shape[:2]
+    result = {"6nm": np.empty((0, 2), dtype=np.float64),
+              "12nm": np.empty((0, 2), dtype=np.float64)}
+
+    for cls, paths in [("6nm", record.csv_6nm_paths),
+                       ("12nm", record.csv_12nm_paths)]:
+        all_coords = []
+        for csv_path in paths:
+            df = load_annotations_csv(csv_path)
+            # Convert microns to pixels
+            px_x = df["X"].values * MICRONS_TO_PIXELS
+            px_y = df["Y"].values * MICRONS_TO_PIXELS
+            # Validate: coords must fall within image bounds
+            assert px_x.max() < w + 10, \
+                f"X coords out of bounds ({px_x.max():.0f} > {w}) in {csv_path}"
+            assert px_y.max() < h + 10, \
+                f"Y coords out of bounds ({px_y.max():.0f} > {h}) in {csv_path}"
+            all_coords.append(np.stack([px_x, px_y], axis=1))
+
+        if all_coords:
+            coords = np.concatenate(all_coords, axis=0)
+            # Deduplicate (for S22 merged results): remove within 3px
+            if len(coords) > 1:
+                coords = _deduplicate_coords(coords, min_dist=3.0)
+            result[cls] = coords
+
+    return result
+
+
+def _deduplicate_coords(
+    coords: np.ndarray, min_dist: float = 3.0
+) -> np.ndarray:
+    """Remove duplicate coordinates within min_dist pixels."""
+    from scipy.spatial.distance import cdist
+
+    if len(coords) <= 1:
+        return coords
+    dists = cdist(coords, coords)
+    np.fill_diagonal(dists, np.inf)
+    keep = np.ones(len(coords), dtype=bool)
+    for i in range(len(coords)):
+        if not keep[i]:
+            continue
+        # Mark later duplicates
+        for j in range(i + 1, len(coords)):
+            if keep[j] and dists[i, j] < min_dist:
+                keep[j] = False
+    return coords[keep]
+
+
+# ---------------------------------------------------------------------------
+# Preprocessing transforms
+# ---------------------------------------------------------------------------
+
+def preprocess_image(img: np.ndarray, bead_class: str,
+                     tophat_radii: Optional[Dict[str, int]] = None,
+                     clahe_clip_limit: float = 0.03,
+                     clahe_kernel_size: int = 64) -> np.ndarray:
+    """
+    Top-hat + CLAHE preprocessing. Used ONLY by LodeStar (Stage 1).
+
+    Not used for model training — the CEM500K backbone expects raw EM images.
+    """
+    from skimage import exposure
+    from skimage.morphology import disk, white_tophat
+
+    if tophat_radii is None:
+        tophat_radii = {"6nm": 8, "12nm": 12}
+
+    img_inv = (255 - img).astype(np.float32)
+    radius = tophat_radii[bead_class]
+    tophat = white_tophat(img_inv, disk(radius))
+
+    tophat_max = tophat.max()
+    if tophat_max > 0:
+        tophat_norm = tophat / tophat_max
+    else:
+        tophat_norm = tophat
+
+    enhanced = exposure.equalize_adapthist(
+        tophat_norm,
+        clip_limit=clahe_clip_limit,
+        kernel_size=clahe_kernel_size,
+    )
+    return (enhanced * 255).astype(np.uint8)
+
+
+# ---------------------------------------------------------------------------
+# Convenience: load everything for one synapse
+# ---------------------------------------------------------------------------
+
+def load_synapse(record: SynapseRecord) -> dict:
+    """
+    Load image, mask, and annotations for one synapse.
+
+    Returns dict with keys: 'image', 'mask', 'annotations',
+                            'synapse_id', 'image_shape'
+    """
+    img = load_image(record.image_path)
+    mask = load_mask(record.mask_path) if record.mask_path else None
+    annotations = load_all_annotations(record, img.shape)
+
+    return {
+        "synapse_id": record.synapse_id,
+        "image": img,
+        "mask": mask,
+        "annotations": annotations,
+        "image_shape": img.shape,
+    }