diff.py

from functools import cache
from pathlib import Path

import cv2
import numpy as np
import plotly.graph_objects as go
import torch
from funcy import print_durations
from huggingface_hub import hf_hub_download
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForDepthEstimation, pipeline

from depth_anything_v2.dpt import DepthAnythingV2

DEVICE = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using device: {DEVICE}")


@cache
def model_original(encoder="vits"):
    model_configs = {
        "vits": {
            "encoder": "vits",
            "features": 64,
            "out_channels": [48, 96, 192, 384],
        },
        "vitb": {
            "encoder": "vitb",
            "features": 128,
            "out_channels": [96, 192, 384, 768],
        },
        "vitl": {
            "encoder": "vitl",
            "features": 256,
            "out_channels": [256, 512, 1024, 1024],
        },
        "vitg": {
            "encoder": "vitg",
            "features": 384,
            "out_channels": [1536, 1536, 1536, 1536],
        },
    }
    checkpoints = {
        "vits": {
            "repo_id": "depth-anything/Depth-Anything-V2-Small",
            "filename": "depth_anything_v2_vits.pth",
        },
        "vitb": {
            "repo_id": "depth-anything/Depth-Anything-V2-Base",
            "filename": "depth_anything_v2_vitb.pth",
        },
        "vitl": {
            "repo_id": "depth-anything/Depth-Anything-V2-Large",
            "filename": "depth_anything_v2_vitl.pth",
        },
    }

    model = DepthAnythingV2(**model_configs[encoder])
    model.load_state_dict(
        torch.load(
            hf_hub_download(**checkpoints[encoder]),
            map_location="cpu",
        )
    )
    model = model.to(DEVICE).eval()

    return model


@print_durations
def depth_original(
    filename,
    encoder="vits",
    input_size=518,
):
    raw_image = cv2.imread(filename)

    model = model_original(encoder)

    return model.infer_image(raw_image, input_size)


@cache
def model_tf(encoder="vits"):
    model = {
        "vits": "depth-anything/Depth-Anything-V2-Small-hf",
        "vitb": "depth-anything/Depth-Anything-V2-Base-hf",
        "vitl": "depth-anything/Depth-Anything-V2-Large-hf",
    }

    pipe = pipeline(task="depth-estimation", model=model[encoder])
    pipe.model.to(DEVICE)

    return pipe


@print_durations
def depth_tf(
    filename,
    encoder="vits",
):
    image = Image.open(filename)

    pipe = model_tf(encoder)

    depth = pipe(image)["predicted_depth"]

    return depth.cpu().numpy()


@cache
def model_tf_manual(encoder="vits"):
    model = {
        "vits": "depth-anything/Depth-Anything-V2-Small-hf",
        "vitb": "depth-anything/Depth-Anything-V2-Base-hf",
        "vitl": "depth-anything/Depth-Anything-V2-Large-hf",
    }

    image_processor = AutoImageProcessor.from_pretrained(model[encoder])
    model = AutoModelForDepthEstimation.from_pretrained(model[encoder])
    model.to(DEVICE)

    return image_processor, model


# Ref: transformers.DPTImageProcessor.post_process_depth_estimation
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/dpt/image_processing_dpt.py#L629-L668
# Converted from method into function, removed type hints and `requires_backends` check
def post_process_depth_estimation(
    outputs,
    target_sizes=None,
):
    """
    Converts the raw output of [`DepthEstimatorOutput`] into final depth predictions and depth PIL images.
    Only supports PyTorch.

    Args:
        outputs ([`DepthEstimatorOutput`]):
            Raw outputs of the model.
        target_sizes (`TensorType` or `List[Tuple[int, int]]`, *optional*):
            Tensor of shape `(batch_size, 2)` or list of tuples (`Tuple[int, int]`) containing the target size
            (height, width) of each image in the batch. If left to None, predictions will not be resized.

    Returns:
        `List[Dict[str, TensorType]]`: A list of dictionaries of tensors representing the processed depth
        predictions.
    """

    predicted_depth = outputs.predicted_depth

    if (target_sizes is not None) and (len(predicted_depth) != len(target_sizes)):
        raise ValueError(
            "Make sure that you pass in as many target sizes as the batch dimension of the predicted depth"
        )

    results = []
    target_sizes = (
        [None] * len(predicted_depth) if target_sizes is None else target_sizes
    )
    for depth, target_size in zip(predicted_depth, target_sizes):
        if target_size is not None:
            depth = torch.nn.functional.interpolate(
                depth.unsqueeze(0).unsqueeze(1),
                size=target_size,
                mode="bilinear",
                align_corners=True,
            ).squeeze()

        results.append({"predicted_depth": depth})

    return results


@print_durations
def depth_tf_custom_pp(
    filename,
    encoder="vits",
):
    image = Image.open(filename)

    image_processor, model = model_tf_manual(encoder)

    inputs = image_processor(image, return_tensors="pt").to(DEVICE)
    with torch.inference_mode():
        outputs = model(**inputs)

    post_processed_output = post_process_depth_estimation(
        outputs,
        target_sizes=[(image.height, image.width)],
    )

    return post_processed_output[0]["predicted_depth"].detach().cpu().numpy()


@print_durations
def depth_tf_preprocess(
    filename,
    encoder="vits",
    input_size=518,
):
    raw_image = cv2.imread(filename)

    _, model = model_tf_manual(encoder)

    pixel_values, (h, w) = DepthAnythingV2.image2tensor(
        raw_image, input_size=input_size
    )

    with torch.inference_mode():
        output = model(pixel_values=pixel_values)

    depth = DepthAnythingV2.interpolate_depth(output["predicted_depth"], h, w)

    return depth.detach().cpu().numpy()


def plot_heightmap_3d(height_data, title="3D Heightmap", colorscale="viridis"):
    h, w = height_data.shape
    x = np.arange(w)
    y = np.arange(h)

    fig = go.Figure(
        data=[
            go.Surface(
                z=height_data,
                x=x,
                y=y,
                colorscale=colorscale,
                hovertemplate="X: %{x}<br>Y: %{y}<br>Height: %{z:.3f}<extra></extra>",
                colorbar=dict(title="Height Value"),
            )
        ]
    )

    fig.update_layout(
        title=title,
        scene=dict(
            xaxis_title="Width (pixels)",
            yaxis_title="Height (pixels)",
            zaxis_title="Depth/Height Value",
            # camera=dict(eye=dict(x=1.5, y=1.5, z=1.5)),
            aspectmode="manual",
            # aspectratio=dict(x=1, y=1, z=0.5),
        ),
    )

    return fig


encoder = "vits"
input_size = 518
assets = Path(__file__).parent / "assets" / "examples"
output = Path(__file__).parent / "depth_vis"
output.mkdir(parents=True, exist_ok=True)

for file in assets.glob("*.jpg"):
    output_dir = output / file.name
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Processing {file}")
    o = depth_original(file, encoder=encoder, input_size=input_size)
    t = depth_tf(file, encoder=encoder)
    tc = depth_tf_custom_pp(file, encoder=encoder)
    tp = depth_tf_preprocess(file, encoder=encoder, input_size=input_size)

    np.save(output_dir / "o.npy", o)
    np.save(output_dir / "t.npy", t)
    np.save(output_dir / "tc.npy", tc)
    np.save(output_dir / "tp.npy", tp)

    print("  Differences between original and transformers:")
    d = np.abs(o - t)
    print(f"    Sum of absolute differences: {np.sum(d)}")
    print(f"    Max diff: {d.max()}")
    print(f"    Mean diff: {d.mean()}")

    # plot_heightmap_3d(d).show()

    print("  Differences between original and tf custom post-processor:")
    d = np.abs(o - tc)
    print(f"    Sum of absolute differences: {np.sum(d)}")
    print(f"    Max diff: {d.max()}")
    print(f"    Mean diff: {d.mean()}")

    plot_heightmap_3d(d).show()

    # print("  Differences between transformers and tf+original preprocessor:")
    # d = np.abs(t - tp)
    # print(f"    Sum of absolute differences: {np.sum(d)}")
    # print(f"    Max diff: {d.max()}")
    # print(f"    Mean diff: {d.mean()}")

    print("  Differences between original and tf+original preprocessor:")
    d = np.abs(o - tp)
    print(f"    Sum of absolute differences: {np.sum(d)}")
    print(f"    Max diff: {d.max()}")
    print(f"    Mean diff: {d.mean()}")

    # plot_heightmap_3d(d).show()

    break