Depth Anything V2 - slgrobotics/articubot_one GitHub Wiki

#!/usr/bin/env python3

import time
import cv2
import numpy as np
import torch
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForDepthEstimation

MODEL_NAME = "depth-anything/Depth-Anything-V2-Metric-Indoor-Base-hf"
#MODEL_NAME = "depth-anything/Depth-Anything-V2-Metric-Outdoor-Base-hf"
#MODEL_NAME = "depth-anything/Depth-Anything-V2-Base-hf"

print(f"MODEL_NAME: {MODEL_NAME}")

device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Device: {device}")

if device == "cuda":
    print(f"GPU: {torch.cuda.get_device_name(0)}")

print("Loading model...")

processor = AutoImageProcessor.from_pretrained(MODEL_NAME)

print("==================================================")
print("Processor size    :", processor.size)
print("keep_aspect_ratio :", processor.keep_aspect_ratio)
print("ensure_multiple_of:", processor.ensure_multiple_of)
print("==================================================")

model = AutoModelForDepthEstimation.from_pretrained(
    MODEL_NAME
).to(device)

model.eval()

print("Model loaded.")

image = Image.open("test.png").convert("RGB")
#image = Image.open("test.jpeg").convert("RGB")

orig_width, orig_height = image.size
print(f"Original shape: {orig_width}x{orig_height}")

inputs = processor(images=image, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}


# Warm up GPU
with torch.no_grad():
    _ = model(**inputs)

torch.cuda.synchronize() if device == "cuda" else None

start = time.perf_counter()

with torch.no_grad():
    outputs = model(**inputs)

torch.cuda.synchronize() if device == "cuda" else None

elapsed = time.perf_counter() - start

# if not resizing:
#depth = outputs.predicted_depth.squeeze().cpu().numpy()

# Keep it as a torch Tensor for resizing
depth_tensor = outputs.predicted_depth.unsqueeze(1)

print(f"Inference time: {elapsed * 1000:.1f} ms")
print(f"Depth shape: {depth_tensor.shape}")
print(f"Minimum depth: {depth_tensor.min():.2f} m")
print(f"Maximum depth: {depth_tensor.max():.2f} m")

print(f"Resized depth shape: {depth_tensor.shape}")

depth_tensor = torch.nn.functional.interpolate(
    depth_tensor,
    size=(orig_height, orig_width),
    mode="bicubic",
    align_corners=False,
)

# Convert to NumPy only after resizing
depth = depth_tensor.squeeze().cpu().numpy()
print(f"Center depth: {depth[depth.shape[0] // 2, depth.shape[1] // 2]:.2f} m")

# Exact floating-point model result
np.save("depth.npy", depth)
print("Saved depth.npy")

# Machine-readable 16-bit depth image, millimeters (16UC1 - ROS convention)
depth_mm = np.round(depth * 1000.0).astype(np.uint16)
cv2.imwrite("depth_16.png", depth_mm)
print("Saved depth_16.png")

# Produce a visualization.
depth_vis = depth.copy()

depth_vis -= depth_vis.min()

if depth_vis.max() > 0:
    depth_vis /= depth_vis.max()

depth_vis = (depth_vis * 255).astype(np.uint8)
depth_vis = cv2.applyColorMap(depth_vis, cv2.COLORMAP_INFERNO)

cv2.imwrite("depth.jpg", depth_vis)

print("Saved depth.jpg")

#!/usr/bin/env python3

import argparse
from pathlib import Path

import cv2
import numpy as np


PANEL_WIDTH = 320


def load_depth(filename):
    """
    Load a metric depth map.

    Supported formats:
      .npy  - floating-point depth in meters
      .png  - 16-bit unsigned depth in millimeters

    Returns:
      depth      - float32 NumPy array containing depth in meters
      raw_depth  - original stored values
      source     - description of the source format
    """

    path = Path(filename)
    extension = path.suffix.lower()

    if extension == ".npy":
        raw_depth = np.load(filename)

        if raw_depth.ndim != 2:
            raise RuntimeError(
                f"Expected a 2-D depth array, got shape {raw_depth.shape}"
            )

        depth = raw_depth.astype(np.float32)

        source = "NumPy float depth (meters)"

    elif extension == ".png":
        raw_depth = cv2.imread(filename, cv2.IMREAD_UNCHANGED)

        if raw_depth is None:
            raise RuntimeError(f"Unable to read {filename}")

        if raw_depth.ndim != 2:
            raise RuntimeError(
                f"Expected a single-channel depth PNG, "
                f"got shape {raw_depth.shape}"
            )

        if raw_depth.dtype != np.uint16:
            raise RuntimeError(
                f"Expected a 16-bit depth PNG, got {raw_depth.dtype}"
            )

        # PNG stores depth in millimeters.
        depth = raw_depth.astype(np.float32) / 1000.0

        source = "16-bit PNG depth (millimeters)"

    else:
        raise RuntimeError(
            f"Unsupported depth file format: {extension}\n"
            "Supported formats are .npy and 16-bit .png"
        )

    return depth, raw_depth, source


def create_depth_visualization(depth):
    """
    Convert the floating-point depth map to a color visualization.
    """

    valid = np.isfinite(depth) & (depth > 0)

    if not np.any(valid):
        raise RuntimeError("Depth image contains no valid values.")

    min_depth = np.min(depth[valid])
    max_depth = np.max(depth[valid])

    depth_norm = np.zeros_like(depth, dtype=np.float32)

    if max_depth > min_depth:
        depth_norm[valid] = (
            (depth[valid] - min_depth) /
            (max_depth - min_depth)
        )

    # Near = bright/hot, far = dark.
    depth_norm[valid] = 1.0 - depth_norm[valid]

    depth_u8 = (depth_norm * 255).astype(np.uint8)

    return cv2.applyColorMap(depth_u8, cv2.COLORMAP_INFERNO)


def main():
    parser = argparse.ArgumentParser(
        description="Interactively inspect a Depth Anything V2 depth map."
    )

    parser.add_argument(
        "depth_file",
        nargs="?",
        default="depth.npy",
        help="Depth map (.npy meters or 16-bit .png millimeters)"
    )

    args = parser.parse_args()

    print(f"Loading {args.depth_file}...")

    depth, raw_depth, source = load_depth(args.depth_file)

    height, width = depth.shape

    valid = np.isfinite(depth) & (depth > 0)

    print(f"Source: {source}")
    print(f"Depth image: {width} x {height}")
    print(f"Data type: {raw_depth.dtype}")
    print(f"Minimum depth: {np.min(depth[valid]):.3f} m")
    print(f"Maximum depth: {np.max(depth[valid]):.3f} m")

    depth_vis = create_depth_visualization(depth)

    display = np.zeros(
        (height, width + PANEL_WIDTH, 3),
        dtype=np.uint8
    )

    window_name = "Depth Inspector"

    mouse_x = -1
    mouse_y = -1


    def redraw():
        display[:, :width] = depth_vis

        # Clear information panel.
        display[:, width:] = 30

        panel_x = width + 20

        cv2.putText(
            display,
            "DEPTH INSPECTOR",
            (panel_x, 35),
            cv2.FONT_HERSHEY_SIMPLEX,
            1.0,
            (255, 255, 255),
            2,
            cv2.LINE_AA,
        )

        cv2.putText(
            display,
            f"Image: {width} x {height}",
            (panel_x, 75),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (220, 220, 220),
            1,
            cv2.LINE_AA,
        )

        cv2.putText(
            display,
            f"Range: {np.min(depth[valid]):.2f} - "
            f"{np.max(depth[valid]):.2f} m",
            (panel_x, 100),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (220, 220, 220),
            1,
            cv2.LINE_AA,
        )

        if 0 <= mouse_x < width and 0 <= mouse_y < height:

            z = depth[mouse_y, mouse_x]
            raw = raw_depth[mouse_y, mouse_x]

            b, g, r = depth_vis[mouse_y, mouse_x]

            cv2.putText(
                display,
                f"Pixel: ({mouse_x}, {mouse_y})",
                (panel_x, 155),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                (255, 255, 255),
                1,
                cv2.LINE_AA,
            )

            #
            # Display the actual stored value.
            #
            if raw_depth.dtype == np.uint16:
                raw_text = f"Raw: {int(raw)} mm"
            else:
                raw_text = f"Raw: {float(raw):.6f}"

            cv2.putText(
                display,
                raw_text,
                (panel_x, 195),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.55,
                (220, 220, 220),
                1,
                cv2.LINE_AA,
            )

            cv2.putText(
                display,
                f"RGB: ({r}, {g}, {b})",
                (panel_x, 230),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (180, 180, 180),
                1,
                cv2.LINE_AA,
            )

            if np.isfinite(z) and z > 0:

                cv2.putText(
                    display,
                    f"Depth: {z:.3f} m",
                    (panel_x, 285),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.75,
                    (255, 255, 255),
                    2,
                    cv2.LINE_AA,
                )

                cv2.putText(
                    display,
                    f"{z * 100:.1f} cm",
                    (panel_x, 320),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.6,
                    (220, 220, 220),
                    1,
                    cv2.LINE_AA,
                )

            else:

                cv2.putText(
                    display,
                    "Depth: INVALID",
                    (panel_x, 285),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.7,
                    (255, 255, 255),
                    2,
                    cv2.LINE_AA,
                )

            cv2.drawMarker(
                display,
                (mouse_x, mouse_y),
                (255, 255, 255),
                markerType=cv2.MARKER_CROSS,
                markerSize=20,
                thickness=1,
            )

        cv2.putText(
            display,
            "Move mouse over image",
            (panel_x, height - 65),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.45,
            (180, 180, 180),
            1,
            cv2.LINE_AA,
        )

        cv2.putText(
            display,
            "ESC or Q to exit",
            (panel_x, height - 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            1.0,
            (180, 180, 180),
            2,
            cv2.LINE_AA,
        )


    def mouse_callback(event, x, y, flags, param):
        nonlocal mouse_x, mouse_y

        if event == cv2.EVENT_MOUSEMOVE:
            mouse_x = x
            mouse_y = y
            redraw()


    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
    cv2.setMouseCallback(window_name, mouse_callback)

    redraw()

    while True:
        cv2.imshow(window_name, display)

        # Set initial window size
        cv2.resizeWindow(window_name, 1400, 900)

        key = cv2.waitKey(20) & 0xFF

        if key == 27 or key == ord("q"):
            break

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()