import argparse from pathlib import Path import cv2 import numpy as np import torch from PIL import Image from transformers import pipeline IMG_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"} def robust_normalize(depth: np.ndarray, invert: bool = False) -> np.ndarray: d = depth.astype(np.float32) p2 = np.percentile(d, 2) p98 = np.percentile(d, 98) dn = (d - p2) / (p98 - p2 + 1e-6) dn = np.clip(dn, 0.0, 1.0) if invert: dn = 1.0 - dn return dn def depth_to_colormap(depth_norm: np.ndarray) -> np.ndarray: u8 = (depth_norm * 255).astype(np.uint8) return cv2.applyColorMap(u8, cv2.COLORMAP_TURBO) def make_bands(depth_norm: np.ndarray) -> np.ndarray: bands = np.zeros_like(depth_norm, dtype=np.uint8) bands[depth_norm >= 0.33] = 1 bands[depth_norm >= 0.66] = 2 out = np.zeros((bands.shape[0], bands.shape[1], 3), dtype=np.uint8) # BGR out[bands == 0] = (80, 80, 255) # faixa 0 out[bands == 1] = (80, 255, 255) # faixa 1 out[bands == 2] = (80, 255, 80) # faixa 2 return out def run_depth(pipe, image_path: Path, invert: bool): img_pil = Image.open(image_path).convert("RGB") rgb = np.array(img_pil) preview_bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) result = pipe(img_pil) if "predicted_depth" in result: depth = result["predicted_depth"] if hasattr(depth, "detach"): depth = depth.detach().cpu().numpy() depth = np.array(depth).squeeze().astype(np.float32) else: depth_img = result["depth"] depth = np.array(depth_img).astype(np.float32) if depth.ndim == 3: depth = cv2.cvtColor( depth.astype(np.uint8), cv2.COLOR_RGB2GRAY ).astype(np.float32) depth = cv2.resize( depth, (preview_bgr.shape[1], preview_bgr.shape[0]), interpolation=cv2.INTER_CUBIC ) depth_norm = robust_normalize(depth, invert=invert) depth_color = depth_to_colormap(depth_norm) depth_bands = make_bands(depth_norm) return preview_bgr, depth_norm, depth_color, depth_bands def resize_to_height(img: np.ndarray, target_h: int) -> np.ndarray: h, w = img.shape[:2] if h == target_h: return img scale = target_h / h new_w = int(w * scale) return cv2.resize(img, (new_w, target_h), interpolation=cv2.INTER_AREA) def compose_view(preview_bgr, depth_color, depth_bands, image_path, index, total, mode): target_h = 520 left = resize_to_height(preview_bgr, target_h) if mode == "depth": right_img = depth_color right_title = "Depth Anything V2" else: right_img = depth_bands right_title = "Depth bands" right = resize_to_height(right_img, target_h) # Garante mesma altura h = min(left.shape[0], right.shape[0]) left = left[:h] right = right[:h] canvas = np.hstack([left, right]) text1 = f"{index + 1}/{total} - {image_path.name}" text2 = f"Modo: {right_title} | N/SPACE prox | A ant | M modo | I invert | S salvar | Q sair" cv2.rectangle(canvas, (0, 0), (canvas.shape[1], 58), (0, 0, 0), -1) cv2.putText(canvas, text1, (12, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.58, (255, 255, 255), 1, cv2.LINE_AA) cv2.putText(canvas, text2, (12, 48), cv2.FONT_HERSHEY_SIMPLEX, 0.50, (220, 220, 220), 1, cv2.LINE_AA) return canvas def main(): parser = argparse.ArgumentParser() parser.add_argument("--input_dir", required=True, help="Pasta com imagens preview PNG/JPG") parser.add_argument("--model", default="depth-anything/Depth-Anything-V2-Small-hf") parser.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"]) parser.add_argument("--invert", action="store_true") parser.add_argument("--save_dir", default="depth_viewer_saves") args = parser.parse_args() input_dir = Path(args.input_dir) save_dir = Path(args.save_dir) save_dir.mkdir(parents=True, exist_ok=True) image_paths = sorted([ p for p in input_dir.rglob("*") if p.suffix.lower() in IMG_EXTS ]) if not image_paths: raise RuntimeError(f"Nenhuma imagem encontrada em: {input_dir}") if args.device == "auto": device = 0 if torch.cuda.is_available() else -1 elif args.device == "cuda": device = 0 else: device = -1 print(f"[INFO] imagens: {len(image_paths)}") print(f"[INFO] modelo: {args.model}") print(f"[INFO] device: {'cuda' if device == 0 else 'cpu'}") print("[INFO] controles:") print(" N ou SPACE = próxima") print(" A = anterior") print(" M = alterna depth/faixas") print(" I = inverte depth") print(" S = salva visual atual") print(" Q ou ESC = sair") pipe = pipeline( task="depth-estimation", model=args.model, device=device ) idx = 0 invert = args.invert mode = "depth" cached_path = None cached_data = None cv2.namedWindow("Depth Anything V2 Viewer", cv2.WINDOW_NORMAL) while True: image_path = image_paths[idx] need_reprocess = cached_path != image_path or cached_data is None if need_reprocess: print(f"[RUN] {idx + 1}/{len(image_paths)} - {image_path.name}") preview_bgr, depth_norm, depth_color, depth_bands = run_depth(pipe, image_path, invert=invert) cached_data = (preview_bgr, depth_norm, depth_color, depth_bands) cached_path = image_path else: preview_bgr, depth_norm, depth_color, depth_bands = cached_data view = compose_view( preview_bgr=preview_bgr, depth_color=depth_color, depth_bands=depth_bands, image_path=image_path, index=idx, total=len(image_paths), mode=mode ) cv2.imshow("Depth Anything V2 Viewer", view) key = cv2.waitKey(0) & 0xFF if key in [27, ord("q"), ord("Q")]: break elif key in [ord("n"), ord("N"), 32]: idx = min(idx + 1, len(image_paths) - 1) cached_path = None elif key in [ord("a"), ord("A")]: idx = max(idx - 1, 0) cached_path = None elif key in [ord("m"), ord("M")]: mode = "bands" if mode == "depth" else "depth" elif key in [ord("i"), ord("I")]: invert = not invert cached_path = None print(f"[INFO] invert={invert}") elif key in [ord("s"), ord("S")]: out_path = save_dir / f"{image_path.stem}_viewer_{mode}.png" cv2.imwrite(str(out_path), view) print(f"[SAVE] {out_path}") cv2.destroyAllWindows() if __name__ == "__main__": main()