Ininitial commit

app.py

@@ -0,0 +1,24 @@
+from depth import MidasDepth
+import gradio as gr
+import numpy as np
+import cv2
+
+
+depth_estimator = MidasDepth()
+
+
+def get_depth(rgb):
+    depth = depth_estimator.get_depth(rgb)
+
+    return rgb, (depth.clip(0, 64) * 1024).astype("uint16")
+
+
+starter = gr.Interface(fn=get_depth, inputs=[
+    gr.components.Image(label="rgb", type="pil"),
+], outputs=[
+    gr.components.Image(type="pil", label="image"),
+    gr.components.Image(type="numpy", label="depth"),
+
+])
+
+gr.Interface(get_depth).launch(share=True)

depth.py

@@ -0,0 +1,60 @@
+from typing import Optional, Union
+from tqdm.auto import trange
+from PIL import ImageOps
+from PIL import Image
+from torch import nn
+import numpy as np
+import torch
+import cv2
+
+
+class MidasDepth(nn.Module):
+    def __init__(self, model_type="DPT_Large",
+                 device=torch.device(
+                     "cuda" if torch.cuda.is_available() else "cpu"),
+                 is_inpainting=False):
+        super().__init__()
+        self.device = device
+        if self.device.type == "mps":
+            self.device = torch.device("cpu")
+        self.model = torch.hub.load(
+            "intel-isl/MiDaS", model_type).to(self.device).eval().requires_grad_(False)
+        self.transform = torch.hub.load(
+            "intel-isl/MiDaS", "transforms").dpt_transform
+
+    @torch.no_grad()
+    def forward(self, image):
+        if torch.is_tensor(image):
+            image = image.cpu().detach()
+        if not isinstance(image, np.ndarray):
+            image = np.asarray(image)
+        image = image.squeeze()
+        batch = self.transform(image).to(self.device)
+        prediction = self.model(batch)
+        prediction = torch.nn.functional.interpolate(
+            prediction.unsqueeze(1),
+            size=image.shape[-3:-1],
+            mode="bicubic",
+            align_corners=False,
+        )[:, 0]
+        # prediction = prediction - prediction.min() + 1.5
+        # prediction = 20 / prediction
+        return prediction  # .squeeze()
+
+    @torch.no_grad()
+    def get_depth(self, img):
+        im = torch.from_numpy(np.asarray(img)).float().to(self.device) / 255.
+        og_depth = self(im.unsqueeze(0) * 255.)[0]
+        d = og_depth
+        d = (d - d.min()) / (d.max() - d.min()) * (10 - 3) + 3
+        d = 30 / d
+        # d = d.max() - d
+        # d = d / d.max() * 15
+        # d = d + 1.5
+        return d.detach().cpu().numpy()
+
+
+if __name__ == "__main__":
+    from matplotlib import pyplot as plt
+    plt.imshow(MidasDepth().get_depth(Image.open("horse.jpg")))
+    plt.show()

serve_modal.py

@@ -0,0 +1,9 @@
+import modal
+
+
+stub = modal.Stub()
+
+
+@stub.function()
+def estimate_depth(image):
+    pass