Update app.py

app.py

@@ -1,40 +1,45 @@
 import os
 import gradio as gr
-from transformers import pipeline
-
+import torch
 import torch.nn.functional as F
+import numpy as np
 from PIL import Image
+from transformers import pipeline
 
-depth_estimator = pipeline(task="depth-estimation",
-                        model="Intel/dpt-hybrid-midas")
+depth_estimator = pipeline(task="depth-estimation", model="Intel/dpt-hybrid-midas")
 
 def launch(input_image):
     out = depth_estimator(input_image)
 
-    # resize the prediction
+    predicted_depth = torch.tensor(out["predicted_depth"])
+    
+    if len(predicted_depth.shape) == 2:  # Если двумерен, добавляем оси
+        predicted_depth = predicted_depth.unsqueeze(0).unsqueeze(0)
+    
     prediction = F.interpolate(
-        out["predicted_depth"].unsqueeze(1),
-        size=input_image.size[::-1],
+        predicted_depth,
+        size=input_image.size[::-1],  # Порядок: (ширина, высота)
         mode="bicubic",
         align_corners=False,
     )
 
-    # normalize the prediction
     output = prediction.squeeze().numpy()
     formatted = (output * 255 / np.max(output)).astype("uint8")
     depth = Image.fromarray(formatted)
     return depth
 
-iface = gr.Interface(launch,
-                     inputs=gr.Image(type='pil'),
-                     outputs=gr.Image(type='pil'))
+iface = gr.Interface(
+    launch,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+)
 
 demo = gr.Blocks()
 
 with demo:
     gr.TabbedInterface(
         [iface],
-        ["iface"],
+        ["Depth Estimation Interface"],
     )
 
-demo.launch(debug=True)
+demo.launch(debug=True)