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,36 +191,34 @@ 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",
-                colour="cyan",
-                bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]")
     processed = 0
-    for start in pbar:
+    for (H, W), group in by_size.items():
-        chunk = valid[start : start + batch_size]
+        label = f"🔪 edges {H}x{W} (sobel)"
-        depth_tensors = []
+        pbar = tqdm(range(0, len(group), batch_size), desc=label, unit="batch",
-        for r in chunk:
+                    colour="cyan",
-            np_p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
+                    bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]")
-            vis_p = vis_path(r.output_root, "depth", r.rel_parent, r.stem)
+        for start in pbar:
-            if np_p.exists():
+            chunk = group[start : start + batch_size]
-                d = np.load(np_p).astype(np.float32)
+            depth_tensors = [_load_depth_tensor(r) for r in chunk]
-            else:
+            depths = torch.stack(depth_tensors)
-                from PIL import Image
+            if depths.ndim == 3:
-                d = np.array(Image.open(vis_p)).astype(np.float32) / 255.0
+                depths = depths.unsqueeze(1)
-                if d.ndim == 2:
+            edges_batch = compute_edges_from_depth(depths)
-                    d = d[np.newaxis]
+            for j, r in enumerate(chunk):
-            depth_tensors.append(torch.from_numpy(d))
+                save_edges_async(edges_batch[j], r.output_root, r.rel_parent,
-        depths = torch.stack(depth_tensors)
+                                 stem=r.stem,
-        if depths.ndim == 3:
+                                 save_npy=pipeline_conf.save_npy,
-            depths = depths.unsqueeze(1)
+                                 save_vis=pipeline_conf.save_vis)
-        edges_batch = compute_edges_from_depth(depths)
+            processed += len(chunk)
-        for j, r in enumerate(chunk):
+            pbar.set_postfix(images=f"{processed}/{total_images}")
-            save_edges_async(edges_batch[j], r.output_root, r.rel_parent,
-                             stem=r.stem,
-                             save_npy=pipeline_conf.save_npy,
-                             save_vis=pipeline_conf.save_vis)
-        processed += len(chunk)
-        pbar.set_postfix(images=f"{processed}/{total_images}")
 
     shutdown_io_pool()