Fix edges stage: group by resolution before stacking

Mixed-size datasets (e.g. GTA-UAV: drone 512x512 + satellite 256x256)
caused torch.stack error in run_edges_stage. Now groups depth tensors
by spatial resolution and processes each group separately.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/main.py

@@ -159,6 +159,23 @@ def run_depth_stage(
     unload_model(model)
 
 
+def _load_depth_tensor(r: ImageRecord) -> torch.Tensor:
+    """Load depth from npy or png, return [1, H, W] float32."""
+    np_p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
+    vis_p = vis_path(r.output_root, "depth", r.rel_parent, r.stem)
+    if np_p.exists():
+        d = np.load(np_p).astype(np.float32)
+    else:
+        from PIL import Image as _Img
+        d = np.array(_Img.open(vis_p)).astype(np.float32) / 255.0
+        if d.ndim == 2:
+            d = d[np.newaxis]
+    t = torch.from_numpy(d)
+    if t.ndim == 2:
+        t = t.unsqueeze(0)
+    return t
+
+
 def run_edges_stage(
     records: list[ImageRecord],
     pipeline_conf: PipelineConfig,
@@ -174,25 +191,23 @@ def run_edges_stage(
         else:
             logger.warning("⚠️  No depth for %s, skipping edges.", r.rel_path)
 
+    # Group by spatial resolution to avoid stack errors with mixed sizes.
+    by_size: dict[tuple[int, int], list[ImageRecord]] = {}
+    for r in valid:
+        t = _load_depth_tensor(r)
+        hw = (t.shape[-2], t.shape[-1])
+        by_size.setdefault(hw, []).append(r)
+
     total_images = len(valid)
-    pbar = tqdm(range(0, len(valid), batch_size), desc="🔪 edges (sobel)", unit="batch",
+    processed = 0
+    for (H, W), group in by_size.items():
+        label = f"🔪 edges {H}x{W} (sobel)"
+        pbar = tqdm(range(0, len(group), batch_size), desc=label, unit="batch",
                     colour="cyan",
                     bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]")
-    processed = 0
         for start in pbar:
-        chunk = valid[start : start + batch_size]
-        depth_tensors = []
-        for r in chunk:
-            np_p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
-            vis_p = vis_path(r.output_root, "depth", r.rel_parent, r.stem)
-            if np_p.exists():
-                d = np.load(np_p).astype(np.float32)
-            else:
-                from PIL import Image
-                d = np.array(Image.open(vis_p)).astype(np.float32) / 255.0
-                if d.ndim == 2:
-                    d = d[np.newaxis]
-            depth_tensors.append(torch.from_numpy(d))
+            chunk = group[start : start + batch_size]
+            depth_tensors = [_load_depth_tensor(r) for r in chunk]
             depths = torch.stack(depth_tensors)
             if depths.ndim == 3:
                 depths = depths.unsqueeze(1)