diff --git a/src/main.py b/src/main.py index 1636b8b..d739fc2 100644 --- a/src/main.py +++ b/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: - 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)) - depths = torch.stack(depth_tensors) - if depths.ndim == 3: - depths = depths.unsqueeze(1) - edges_batch = compute_edges_from_depth(depths) - for j, r in enumerate(chunk): - 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}") + 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}]") + for start in pbar: + 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) + edges_batch = compute_edges_from_depth(depths) + for j, r in enumerate(chunk): + 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()