Hugging Face's logo

Spaces:
RaynWu2002
/
DORNet
Runtime error

App Files Community
1
DORNet
File size: 5,607 Bytes 
68c537d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
 
import gradio as gr
import numpy as np
import torch
import os
import cv2
from PIL import Image
import torchvision.transforms as transforms
from net.dornet import Net
from net.dornet_ddp import Net_ddp

# init
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = Net(tiny_model=False).to(device)
model_ckpt_map = {
    "RGB-D-D": "./checkpoints/RGBDD.pth",
    "TOFDSR": "./checkpoints/TOFDSR.pth"
}

# load model
def load_model(model_type: str):
    global net
    ckpt_path = model_ckpt_map[model_type]
    print(f"Loading weights from: {ckpt_path}")
    if model_type == "RGB-D-D":
        net = Net(tiny_model=False).to(device)
    elif model_type == "TOFDSR":
        net = Net_ddp(tiny_model=False).srn.to(device)
    else:
        raise ValueError(f"Unknown model_type: {model_type}")

    net.load_state_dict(torch.load(ckpt_path, map_location=device))
    net.eval()

load_model("RGB-D-D")


# data process
def preprocess_inputs(rgb_image: Image.Image, lr_depth: Image.Image):
    image = np.array(rgb_image.convert("RGB")).astype(np.float32)
    h, w, _ = image.shape
    lr = np.array(lr_depth.resize((w, h), Image.BICUBIC)).astype(np.float32)
    # Normalize depth
    max_out, min_out = 5000.0, 0.0
    lr = (lr - min_out) / (max_out - min_out)
    # Normalize RGB
    maxx, minn = np.max(image), np.min(image)
    image = (image - minn) / (maxx - minn)
    # To tensor
    data_transform = transforms.Compose([transforms.ToTensor()])
    image = data_transform(image).float()
    lr = data_transform(np.expand_dims(lr, 2)).float()
    # Add batch dimension
    lr = lr.unsqueeze(0).to(device)
    image = image.unsqueeze(0).to(device)
    return image, lr, min_out, max_out


# model inference
@torch.no_grad()
def infer(rgb_image: Image.Image, lr_depth: Image.Image, model_type: str):
    load_model(model_type)  # reset weight

    image, lr, min_out, max_out = preprocess_inputs(rgb_image, lr_depth)

    if model_type == "RGB-D-D":
        out = net(x_query=lr, rgb=image)
    elif model_type == "TOFDSR":
        out, _ = net(x_query=lr, rgb=image)

    pred = out[0, 0] * (max_out - min_out) + min_out
    pred = pred.cpu().numpy().astype(np.uint16)
    # raw
    pred_gray = Image.fromarray(pred)

    # heat
    pred_norm = (pred - np.min(pred)) / (np.max(pred) - np.min(pred)) * 255
    pred_vis = pred_norm.astype(np.uint8)
    pred_heat = cv2.applyColorMap(pred_vis, cv2.COLORMAP_PLASMA)
    pred_heat = cv2.cvtColor(pred_heat, cv2.COLOR_BGR2RGB)
    return pred_gray, Image.fromarray(pred_heat)


# Gradio
# demo = gr.Interface(
#     fn=infer,
#     inputs=[
#         gr.Image(label="RGB Image", type="pil"),
#         gr.Image(label="Low-res Depth", type="pil", image_mode="I"),
#         gr.Dropdown(choices=["RGB-D-D", "TOFDSR"], label="Model Type", value="RGB-D-D")
#     ],
#     outputs=[
#         gr.Image(label="DORNet Output", type="pil", elem_classes=["output-image"]),
#         gr.Image(label="Normalized Output", type="pil", elem_classes=["output-image"])
#     ],
#     examples=[
#         ["examples/RGB-D-D/20200518160957_RGB.jpg", "examples/RGB-D-D/20200518160957_LR_fill_depth.png", "RGB-D-D"],
#         ["examples/TOFDSR/2020_09_08_13_59_59_435_rgb_rgb_crop.png", "examples/TOFDSR/2020_09_08_13_59_59_435_rgb_depth_crop_fill.png", "TOFDSR"],
#     ],
#     allow_flagging="never",
#     title="DORNet: A Degradation Oriented and Regularized Network for Blind Depth Super-Resolution \n CVPR 2025 (Oral Presentation)",
#     css="""
#         .output-image {
#             display: flex;
#             justify-content: center;
#             align-items: center;
#         }
#         .output-image img {
#             margin: auto;
#             display: block;
#         }
#         """
# )
#
# demo.launch(share=True)
Intro = """

## DORNet: A Degradation Oriented and Regularized Network for Blind Depth Super-Resolution

[📄 Paper](https://arxiv.org/pdf/2410.11666) • [💻 Code](https://github.com/yanzq95/DORNet) • [📦 Model](https://huggingface.co/wzxwyx/DORNet/tree/main)

"""

with gr.Blocks(css="""

    .output-image {

        display: flex;

        justify-content: center;

        align-items: center;

    }

    .output-image img {

        margin: auto;

        display: block;

    }

""") as demo:
    gr.Markdown(Intro)

    with gr.Row():
        with gr.Column():
            rgb_input = gr.Image(label="RGB Image", type="pil")
            lr_input = gr.Image(label="Low-res Depth", type="pil", image_mode="I")
        with gr.Column():
            output1 = gr.Image(label="DORNet Output", type="pil", elem_classes=["output-image"])
            output2 = gr.Image(label="Normalized Output", type="pil", elem_classes=["output-image"])

    model_selector = gr.Dropdown(choices=["RGB-D-D", "TOFDSR"], label="Model Type", value="RGB-D-D")
    run_button = gr.Button("Run Inference")

    gr.Examples(
        examples=[
            ["examples/RGB-D-D/20200518160957_RGB.jpg", "examples/RGB-D-D/20200518160957_LR_fill_depth.png", "RGB-D-D"],
            ["examples/TOFDSR/2020_09_08_13_59_59_435_rgb_rgb_crop.png", "examples/TOFDSR/2020_09_08_13_59_59_435_rgb_depth_crop_fill.png", "TOFDSR"],
        ],
        inputs=[rgb_input, lr_input, model_selector],
        outputs=[output1, output2],
        label="Try Examples ↓"
    )

    run_button.click(fn=infer, inputs=[rgb_input, lr_input, model_selector], outputs=[output1, output2])

demo.launch(share=True)