update readme

README.md

@@ -1,5 +1,5 @@
 ---
-title: Omnidata Monocular Surface Normal Dpt Hybrid 384
+title: Omnidata Monocular Depth DPT Hybrid 384
 emoji: 🐠
 colorFrom: green
 colorTo: purple
@@ -9,5 +9,5 @@ app_file: app.py
 pinned: false
 license: cc-by-nc-4.0
 ---
-
+# [Use these models in your code:](https://github.com/alexsax/omnidata_models/tree/main)
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -4,13 +4,14 @@ import torch.nn.functional as F
 from torchvision import transforms
 import PIL
 from PIL import Image
-import os
+import matplotlib.pyplot as plt
+import io
 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 = torch.hub.load('alexsax/omnidata_models', 'depth_dpt_hybrid_384')
     model.to(device)
     model.eval()
 
@@ -21,13 +22,14 @@ def setup_transforms(image_size: int) -> transforms.Compose:
         transforms.Resize(image_size, interpolation=PIL.Image.BILINEAR),
         transforms.CenterCrop(image_size),
         transforms.ToTensor(),
+        transforms.Normalize(mean=0.5, std=0.5)
     ])
 
 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:
+def estimate_depth(input_image: PIL.Image.Image) -> PIL.Image.Image:
     with torch.no_grad():
         img_tensor = trans_totensor(input_image)[:3].unsqueeze(0).to(device)
         
@@ -35,16 +37,28 @@ def estimate_surface_normal(input_image: PIL.Image.Image) -> PIL.Image.Image:
             img_tensor = img_tensor.repeat_interleave(3, 1)
 
         output = model(img_tensor).clamp(min=0, max=1)
-        output_image = transforms.ToPILImage()(output[0])
+        output = F.interpolate(output.unsqueeze(0), (512, 512), mode='bicubic').squeeze(0)
+        output = 1 - output.clamp(0, 1)
+
+        # Convert to colormap
+        plt.figure(figsize=(10, 10))
+        plt.imshow(output[0].cpu().numpy(), cmap='viridis')
+        plt.axis('off')
+        
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+        buf.seek(0)
+        output_image = Image.open(buf)
+        plt.close()
 
     return output_image
 
 iface = gr.Interface(
-    fn=estimate_surface_normal,
+    fn=estimate_depth,
     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.",
+    title="Monocular Depth Estimation: Omnidata DPT-Hybrid",
+    description="Upload an image to estimate monocular depth.",
     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",

requirements.txt

@@ -2,4 +2,5 @@ torch>=1.9.0
 torchvision>=0.10.0
 timm==0.4.12
 pillow
-requests
+requests
+matplotlib