app.py

import gradio as gr
import numpy as np
import torch
from torchvision.transforms import Compose
import cv2
from dpt.models import DPTDepthModel
from dpt.transforms import Resize, NormalizeImage, PrepareForNet
import os

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("device: %s" % device)
default_models = {
    "dpt_hybrid": "weights/dpt_hybrid-midas-501f0c75.pt",
    }
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
net_w = net_h = 384
model = DPTDepthModel(
            path=default_models["dpt_hybrid"],
            backbone="vitb_rn50_384",
            non_negative=True,
            enable_attention_hooks=False,
    )
normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
transform = Compose(
        [
            Resize(
                net_w,
                net_h,
                resize_target=None,
                keep_aspect_ratio=True,
                ensure_multiple_of=32,
                resize_method="minimal",
                image_interpolation_method=cv2.INTER_CUBIC,
            ),
            normalization,
            PrepareForNet(),
        ]
    )

model.eval()
model.to(device)

def write_depth(depth, bits=1, absolute_depth=False):
    """Write depth map to pfm and png file.

    Args:
        path (str): filepath without extension
        depth (array): depth
    """
    # write_pfm(path + ".pfm", depth.astype(np.float32))

    if absolute_depth:
        out = depth
    else:
        depth_min = depth.min()
        depth_max = depth.max()

        max_val = (2 ** (8 * bits)) - 1

        if depth_max - depth_min > np.finfo("float").eps:
            out = max_val * (depth - depth_min) / (depth_max - depth_min)
        else:
            out = np.zeros(depth.shape, dtype=depth.dtype)

    if bits == 1:
        return out.astype("uint8")
    elif bits == 2:
        return out.astype("uint16")
    

def DPT(image):
    img_input = transform({"image": image})["image"]
    # compute
    with torch.no_grad():
        sample = torch.from_numpy(img_input).to(device).unsqueeze(0)

        prediction = model.forward(sample)
        prediction = (
            torch.nn.functional.interpolate(
                prediction.unsqueeze(1),
                size=image.shape[:2],
                mode="bicubic",
                align_corners=False,
            )
            .squeeze()
            .cpu()
            .numpy()
        )

        depth_img = write_depth(prediction, bits=2)
        return depth_img

title = " AISeed AI Application Demo "
description = "# A Demo of Deep Learning for Depth Estimation"
example_list = [["examples/" + example] for example in os.listdir("examples")]

with gr.Blocks() as demo:
    demo.title = title
    gr.Markdown(description)
    with gr.Row():
        im = gr.Image(label="Input Image")
        im_2 = gr.Image(label="Depth Image")
        with gr.Column():

            btn1 = gr.Button(value="Depth Estimator")
            btn1.click(DPT, inputs=[im], outputs=[im_2])
            gr.Examples(examples=example_list,
                inputs=[im],
                outputs=[im_2],
                fn=DPT)
    
if __name__ == "__main__":
    demo.launch()