initial commit

app.py

@@ -0,0 +1,63 @@
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+import PIL
+from PIL import Image
+import os
+from typing import Tuple
+
+
+def setup_model(device: torch.device) -> Tuple[torch.nn.Module, int]:
+    image_size = 384
+    model = torch.hub.load('alexsax/omnidata_models', 'surface_normal_dpt_hybrid_384')
+    model.to(device)
+    model.eval()
+
+    return model, image_size
+
+def setup_transforms(image_size: int) -> transforms.Compose:
+    return transforms.Compose([
+        transforms.Resize(image_size, interpolation=PIL.Image.BILINEAR),
+        transforms.CenterCrop(image_size),
+        transforms.ToTensor(),
+    ])
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model, image_size = setup_model(device)
+trans_totensor = setup_transforms(image_size)
+
+def estimate_surface_normal(input_image: PIL.Image.Image) -> PIL.Image.Image:
+    with torch.no_grad():
+        img_tensor = trans_totensor(input_image)[:3].unsqueeze(0).to(device)
+        
+        if img_tensor.shape[1] == 1:
+            img_tensor = img_tensor.repeat_interleave(3, 1)
+
+        output = model(img_tensor).clamp(min=0, max=1)
+        output_image = transforms.ToPILImage()(output[0])
+
+    return output_image
+
+iface = gr.Interface(
+    fn=estimate_surface_normal,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Monocular Surface Normal Estimation: Omnidata DPT-Hybrid",
+    description="Upload an image to estimate monocular surface normals.",
+    examples=[
+        "https://github.com/EPFL-VILAB/omnidata/blob/main/omnidata_tools/torch/assets/test1_rgb.png?raw=true",
+        "https://github.com/EPFL-VILAB/omnidata/blob/main/omnidata_tools/torch/assets/demo/test2.png?raw=true",
+        "https://github.com/EPFL-VILAB/omnidata/blob/main/omnidata_tools/torch/assets/demo/test3.png?raw=true",
+        "https://github.com/EPFL-VILAB/omnidata/blob/main/omnidata_tools/torch/assets/demo/test4.png?raw=true",
+        "https://github.com/EPFL-VILAB/omnidata/blob/main/omnidata_tools/torch/assets/demo/test5.png?raw=true",
+        "https://github.com/EPFL-VILAB/omnidata/blob/main/omnidata_tools/torch/assets/demo/test6.png?raw=true",
+        "https://github.com/EPFL-VILAB/omnidata/blob/main/omnidata_tools/torch/assets/demo/test7.png?raw=true",
+        "https://github.com/EPFL-VILAB/omnidata/blob/main/omnidata_tools/torch/assets/demo/test8.png?raw=true",
+        "https://github.com/EPFL-VILAB/omnidata/blob/main/omnidata_tools/torch/assets/demo/test9.png?raw=true",
+        "https://github.com/EPFL-VILAB/omnidata/blob/main/omnidata_tools/torch/assets/demo/test10.png?raw=true",
+    ],
+)
+
+if __name__ == "__main__":
+    iface.launch()