Create cell4_vae_pipeline.py

cell4_vae_pipeline.py

@@ -0,0 +1,325 @@
+"""
+Cell 4: Multi-Scale Geometric Extraction Pipeline
+===================================================
+Run after Cells 1-3. Uses globals from prior cells.
+
+Updated for PatchCrossAttentionClassifier (no Conv3d).
+Defines extraction functions and the MultiScaleExtractor class.
+Does NOT execute anything — Cell 5 uses these.
+"""
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from dataclasses import dataclass, field
+from typing import List, Optional, Tuple, Dict
+import math
+
+
+@dataclass
+class ExtractionConfig:
+    canonical_shape: Tuple[int, int, int] = (8, 16, 16)
+    scales: List[Tuple[int, int, int]] = field(default_factory=lambda: [
+        (32, 64, 64),   # L0: full latent
+        (16, 32, 32),   # L1: regional
+        (8, 16, 16),    # L2: native patch
+        (4, 8, 8),      # L3: fine detail
+    ])
+    overlap: float = 0.5
+    confidence_threshold: float = 0.7
+    min_occupancy: float = 0.005
+    binarize_percentiles: List[float] = field(default_factory=lambda: [75, 90, 95])
+    n_channel_groups: int = 8
+    device: str = 'cuda'
+
+
+@dataclass
+class GeometricAnnotation:
+    class_name: str
+    class_idx: int
+    confidence: float
+    scale_level: int
+    location: Tuple[int, int, int]
+    patch_size: Tuple[int, int, int]
+    dimension: int = -1
+    is_curved: bool = False
+    curvature_type: str = "none"
+    channel_pair: Optional[Tuple[int, int]] = None
+    channel_group_pair: Optional[Tuple[int, int]] = None
+
+
+def extract_patches_sliding(volume, patch_size, overlap=0.5):
+    """Extract overlapping patches from a 3D volume."""
+    D, H, W = volume.shape
+    pz, py, px = patch_size
+
+    if D < pz or H < py or W < px:
+        pad_d = max(pz - D, 0)
+        pad_h = max(py - H, 0)
+        pad_w = max(px - W, 0)
+        volume = F.pad(volume, (0, pad_w, 0, pad_h, 0, pad_d))
+        D, H, W = volume.shape
+
+    stride_z = max(1, int(pz * (1 - overlap)))
+    stride_y = max(1, int(py * (1 - overlap)))
+    stride_x = max(1, int(px * (1 - overlap)))
+
+    patches = []
+    for z in range(0, max(1, D - pz + 1), stride_z):
+        for y in range(0, max(1, H - py + 1), stride_y):
+            for x in range(0, max(1, W - px + 1), stride_x):
+                patch = volume[z:z+pz, y:y+py, x:x+px]
+                patches.append((patch, (z, y, x)))
+
+    return patches
+
+
+def resize_to_canonical(patch, target=(8, 16, 16)):
+    """Resize 3D patch to canonical resolution via trilinear interpolation."""
+    x = patch.unsqueeze(0).unsqueeze(0).float()
+    x = F.interpolate(x, size=target, mode='trilinear', align_corners=False)
+    return x.squeeze(0).squeeze(0)
+
+
+def binarize_continuous(patch, percentiles=[75, 90, 95]):
+    """Binarize continuous patch at multiple percentile thresholds."""
+    flat = patch.flatten()
+    nonzero = flat[flat.abs() > 1e-8]
+    if len(nonzero) < 10:
+        return [torch.zeros_like(patch)] * len(percentiles)
+    thresholds = [torch.quantile(nonzero.abs(), p / 100.0).item() for p in percentiles]
+    return [(patch.abs() >= t).float() for t in thresholds]
+
+
+def cluster_channels(latents, n_groups=8):
+    """
+    Cluster VAE channels by correlation.
+    latents: (N, C, H, W) batch
+    Returns: (groups, corr_matrix)
+    """
+    N, C, H, W = latents.shape
+    flat = latents.reshape(N, C, -1)
+    flat = flat - flat.mean(dim=-1, keepdim=True)
+    norm = flat.norm(dim=-1, keepdim=True).clamp(min=1e-8)
+    flat = flat / norm
+    corr = torch.bmm(flat, flat.transpose(1, 2)).mean(dim=0)
+
+    dist = (1.0 - corr.abs()).cpu().numpy()
+
+    assigned = [-1] * C
+    groups = []
+    remaining = set(range(C))
+
+    for g in range(n_groups):
+        if not remaining:
+            break
+        if g == 0:
+            seed = min(remaining)
+        else:
+            max_dist = -1
+            seed = min(remaining)
+            for ch in remaining:
+                min_d = min(dist[ch, s] for grp in groups for s in grp)
+                if min_d > max_dist:
+                    max_dist = min_d
+                    seed = ch
+
+        group = [seed]
+        remaining.discard(seed)
+        assigned[seed] = g
+
+        target_size = max(1, C // n_groups)
+        candidates = sorted(remaining, key=lambda c: dist[seed, c])
+        for c in candidates[:target_size - 1]:
+            group.append(c)
+            remaining.discard(c)
+            assigned[c] = g
+
+        groups.append(group)
+
+    if remaining:
+        for ch in remaining:
+            best_g = 0
+            best_d = float('inf')
+            for gi, grp in enumerate(groups):
+                d = min(dist[ch, s] for s in grp)
+                if d < best_d:
+                    best_d = d
+                    best_g = gi
+            groups[best_g].append(ch)
+
+    return groups, corr
+
+
+def compute_inter_group_deviances(latent, groups):
+    """
+    Compute (H, W) deviance maps between channel groups.
+    latent: (C, H, W), groups: list of channel index lists
+    Returns: list of ((group_i, group_j), deviance_map)
+    """
+    group_means = torch.stack([latent[grp].mean(dim=0) for grp in groups])
+    n = len(groups)
+    deviances = []
+    for i in range(n):
+        for j in range(i + 1, n):
+            dev = (group_means[i] - group_means[j]).abs()
+            deviances.append(((i, j), dev))
+    return deviances
+
+
+def deviance_maps_to_3d(deviances, n_groups):
+    """Stack deviance maps into (n_pairs, H, W) volume."""
+    return torch.stack([dev for (_, dev) in deviances], dim=0)
+
+
+class MultiScaleExtractor:
+    """
+    Confidence-cascaded multi-scale geometric extractor.
+    Uses trained PatchCrossAttentionClassifier (from Cell 2 globals).
+    """
+
+    def __init__(self, classifier, config=None):
+        self.classifier = classifier
+        self.config = config or ExtractionConfig()
+        self.classifier.eval()
+
+    @torch.no_grad()
+    def classify_patches(self, patches, max_batch=512):
+        """Classify batch of (B, 8, 16, 16) patches with chunking to avoid OOM."""
+        device = next(self.classifier.parameters()).device
+        N = patches.shape[0]
+
+        all_results = []
+        for start in range(0, N, max_batch):
+            chunk = patches[start:start+max_batch].to(device)
+            out = self.classifier(chunk)
+            probs = F.softmax(out["class_logits"], dim=-1)
+            max_prob, pred_class = probs.max(dim=-1)
+            top2 = probs.topk(2, dim=-1).values
+            margin = top2[:, 0] - top2[:, 1]
+            dim_pred = out["dim_logits"].argmax(dim=-1)
+            curved_pred = (out["is_curved_pred"].squeeze(-1) > 0.0)
+            curv_type_pred = out["curv_type_logits"].argmax(dim=-1)
+
+            all_results.append({
+                "pred_class": pred_class.cpu(),
+                "confidence": margin.cpu(),
+                "max_prob": max_prob.cpu(),
+                "dim_pred": dim_pred.cpu(),
+                "curved_pred": curved_pred.cpu(),
+                "curv_type_pred": curv_type_pred.cpu(),
+                "features": out["features"].cpu(),
+            })
+            del chunk, out, probs
+            torch.cuda.empty_cache()
+
+        return {k: torch.cat([r[k] for r in all_results], dim=0)
+                for k in all_results[0]}
+
+    def extract_from_volume(self, volume, min_confidence=None):
+        """
+        Extract annotations via confidence cascade.
+        volume: (D, H, W) continuous or binary tensor
+        """
+        conf_thresh = min_confidence or self.config.confidence_threshold
+        annotations = []
+
+        regions_to_process = [(0, 0, 0, 0,
+                               volume.shape[0], volume.shape[1], volume.shape[2])]
+
+        for level, scale in enumerate(self.config.scales):
+            if not regions_to_process:
+                break
+
+            next_regions = []
+            pz, py, px = scale
+
+            all_patches = []
+
+            for ridx, (lvl, rz0, ry0, rx0, rz1, ry1, rx1) in enumerate(regions_to_process):
+                if lvl != level:
+                    next_regions.append((lvl, rz0, ry0, rx0, rz1, ry1, rx1))
+                    continue
+
+                subvol = volume[rz0:rz1, ry0:ry1, rx0:rx1]
+                patches = extract_patches_sliding(subvol, (pz, py, px), self.config.overlap)
+
+                for patch, (lz, ly, lx) in patches:
+                    binary_patches = binarize_continuous(patch, self.config.binarize_percentiles)
+                    for bp in binary_patches:
+                        occ = bp.mean().item()
+                        if occ < self.config.min_occupancy:
+                            continue
+                        canonical = resize_to_canonical(bp, self.config.canonical_shape)
+                        all_patches.append((
+                            canonical,
+                            (rz0 + lz, ry0 + ly, rx0 + lx),
+                            ridx, scale))
+
+            if not all_patches:
+                regions_to_process = next_regions
+                continue
+
+            # Batch classify
+            batch = torch.stack([p[0] for p in all_patches])
+            results = self.classify_patches(batch)
+
+            for i, (_, loc, ridx, sc) in enumerate(all_patches):
+                conf = results["confidence"][i].item()
+                cls_idx = results["pred_class"][i].item()
+
+                if conf >= conf_thresh:
+                    ann = GeometricAnnotation(
+                        class_name=CLASS_NAMES[cls_idx],
+                        class_idx=cls_idx,
+                        confidence=conf,
+                        scale_level=level,
+                        location=loc,
+                        patch_size=sc,
+                        dimension=results["dim_pred"][i].item(),
+                        is_curved=bool(results["curved_pred"][i].item()),
+                        curvature_type=CURVATURE_NAMES[results["curv_type_pred"][i].item()],
+                    )
+                    annotations.append(ann)
+                elif level < len(self.config.scales) - 1:
+                    z0, y0, x0 = loc
+                    next_regions.append((
+                        level + 1,
+                        z0, y0, x0,
+                        min(z0 + sc[0], volume.shape[0]),
+                        min(y0 + sc[1], volume.shape[1]),
+                        min(x0 + sc[2], volume.shape[2]),
+                    ))
+
+            regions_to_process = next_regions
+
+        return annotations
+
+    def extract_from_latent(self, latent, channel_groups=None):
+        """
+        Full extraction for a single Flux 2 VAE latent.
+        latent: (C, H, W) tensor
+        """
+        raw_annotations = self.extract_from_volume(latent)
+
+        deviance_annotations = []
+        if channel_groups is not None:
+            deviances = compute_inter_group_deviances(latent, channel_groups)
+            dev_volume = deviance_maps_to_3d(deviances, len(channel_groups))
+            deviance_annotations = self.extract_from_volume(dev_volume)
+            for ann in deviance_annotations:
+                pair_idx = ann.location[0]
+                if pair_idx < len(deviances):
+                    ann.channel_group_pair = deviances[pair_idx][0]
+
+        return {
+            'raw_annotations': raw_annotations,
+            'deviance_annotations': deviance_annotations,
+            'n_raw': len(raw_annotations),
+            'n_deviance': len(deviance_annotations),
+        }
+
+
+print("✓ Cell 4: Extraction pipeline defined (PatchCrossAttention)")
+print(f"  Scales: {ExtractionConfig().scales}")
+print(f"  Canonical: {ExtractionConfig().canonical_shape}")