added depth estimation

app.py

@@ -1,19 +1,63 @@
 import os
 import gradio as gr
+import cv2
+import torch
+import urllib.request
+from PIL import Image
+
+import matplotlib.pyplot as plt
 
 def update(slider, img):
 
     if not os.path.exists('temp'):
       os.system('mkdir temp')
 
-    img.save("temp/image.jpg", "JPEG")
+    filename = "temp/image.jpg"
+
+    img.save(filename, "JPEG")
+
+    model_type = "DPT_Hybrid" 
+    midas = torch.hub.load("intel-isl/MiDaS", model_type)
+    
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    midas.to(device)
+    midas.eval()
+
+    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
+
+    if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
+        transform = midas_transforms.dpt_transform
+    else:
+        transform = midas_transforms.small_transform
+
+    img = cv2.imread(filename)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    input_batch = transform(img).to(device)
+
+    with torch.no_grad():
+        prediction = midas(input_batch)
+
+        prediction = torch.nn.functional.interpolate(
+            prediction.unsqueeze(1),
+            size=img.shape[:2],
+            mode="bicubic",
+            align_corners=False,
+        ).squeeze()
+
+    output = prediction.cpu().numpy()
+
+    out_im = Image.fromarray(output)
+    out_im.save("temp/image_depth.jpg", "JPEG")
+
+    return f'temp/image_depth.jpg'
+
 
-    return f'temp/image.jpg'
 
 with gr.Blocks() as demo:
     gr.Markdown("Start typing below and then click **Run** to see the output.")
     inp = [gr.Slider(1,15, default = 2, label='StepCycles',step= 1)]
-    
+
     with gr.Row():
         inp.append(gr.Image(type="pil", label="Input"))
         out = gr.Image(type="file", label="Output")