main.py

import gradio as gr
from gradio_modal import Modal

import numpy as np

import torch
import torch.nn.functional as F
from torchvision import transforms
import utils.utils as utils

from models.dsine import DSINE
from pathlib import Path

device = torch.device("cpu")

model = DSINE().to(device)
model.pixel_coords = model.pixel_coords.to(device)
model = utils.load_checkpoint("./checkpoints/dsine.pt", model)
model.eval()


def predict_normal(img_np: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
 # normalize
 normalize = transforms.Normalize(
 mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
 )

 with torch.no_grad():
 img = np.array(img_np).astype(np.float32) / 255.0
 img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).to(device)
 _, _, orig_H, orig_W = img.shape

 # zero-pad the input image so that both the width and height are multiples of 32
 l, r, t, b = utils.pad_input(orig_H, orig_W)
 img = F.pad(img, (l, r, t, b), mode="constant", value=0.0)
 img = normalize(img)

 # NOTE: if intrins is not given, we just assume that the principal point is at the center
 # and that the field-of-view is 60 degrees (feel free to modify this assumption)
 intrins = utils.get_intrins_from_fov(
 new_fov=60.0, H=orig_H, W=orig_W, device=device
 ).unsqueeze(0)

 intrins[:, 0, 2] += l
 intrins[:, 1, 2] += t

 pred_norm = model(img, intrins=intrins)[-1]
 pred_norm = pred_norm[:, :, t : t + orig_H, l : l + orig_W]

 # NOTE: by saving the prediction as uint8 png format, you lose a lot of precision
 # if you want to use the predicted normals for downstream tasks, we recommend saving them as float32 NPY files
 pred_norm_np = (
 pred_norm.cpu().detach().numpy()[0, :, :, :].transpose(1, 2, 0)
 ) # (H, W, 3)
 pred_norm_np = ((pred_norm_np + 1.0) / 2.0 * 255.0).astype(np.uint8)

 return pred_norm_np


with gr.Blocks() as demo:
 gr.Markdown(
 """
 # DSINE
 Unofficial Gradio demo of [DSINE: Rethinking Inductive Biases for Surface Normal Estimation](https://github.com/baegwangbin/DSINE)
 """
 )
 with gr.Group():
 with gr.Row():
 input_img = gr.Image(label="Input image", image_mode="RGB")
 output_img = gr.Image(label="Surface Normal")
 # output_img = ImageSlider(label="Surface Normal")

 btn = gr.Button("Predict")
 btn.click(fn=predict_normal, inputs=[input_img], outputs=[output_img])

 example_files = list(Path("examples").glob("*.png"))
 example_files.sort()
 examples = gr.Examples(
 examples=example_files,
 inputs=[input_img],
 outputs=[output_img],
 fn=predict_normal,
 cache_examples=True,
 )

 with Modal(visible=True, allow_user_close=False) as modal:
 gr.Markdown(
 """
 To use this space, you must agree to the terms and conditions.
 Found [HERE](https://github.com/baegwangbin/DSINE/blob/main/LICENSE).
 """,
 )
 btn = gr.Button("I Agree to the Terms and Conditions")
 btn.click(lambda: Modal(visible=False), None, modal)

if __name__ == "__main__":
 demo.launch(server_name="0.0.0.0", server_port=7860)