236 lines
6.7 KiB
Python
236 lines
6.7 KiB
Python
|
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()
|