solve problem with input third try

app.py

@@ -5,28 +5,44 @@ from PIL import Image
 import cv2
 
 # Загрузка модели
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small").to(device)
+midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small")
 midas.eval()
 
-midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms").to(device)
+midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
 transform = midas_transforms.small_transform
 
-
 def predict_depth(image):
-    img = np.array(image)
-    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    input_tensor = transform(img_rgb).to(device)
-
+    # ======= 1. Проверка типа входных данных =======
+    if isinstance(image, torch.Tensor):
+        print(f"Пришёл Tensor с формой: {image.shape}")
+
+        if len(image.shape) == 4:
+            input_tensor = image  # уже батч [1, 3, H, W]
+        elif len(image.shape) == 3:
+            input_tensor = image.unsqueeze(0)  # сделаем батч
+        else:
+            raise ValueError(f"Неожиданный размер Tensor: {image.shape}")
+
+    else:
+        print("Пришёл PIL Image или numpy array")
+        # Если пришло обычное изображение (PIL или numpy)
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+        img = np.array(image)
+        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        input_tensor = transform(img_rgb).unsqueeze(0)
+
+    # ======= 2. Предсказание =======
     with torch.no_grad():
         prediction = midas(input_tensor)
         prediction = torch.nn.functional.interpolate(
             prediction.unsqueeze(1),
-            size=img_rgb.shape[:2],
+            size=(input_tensor.shape[2], input_tensor.shape[3]),
             mode="bicubic",
             align_corners=False,
         ).squeeze()
 
+    # ======= 3. Нормализация карты глубины =======
     depth_map = prediction.cpu().numpy()
     depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
     depth_map = (depth_map * 255).astype(np.uint8)
@@ -34,13 +50,13 @@ def predict_depth(image):
 
     return depth_img
 
-# Интерфейс Gradio
+# Gradio интерфейс
 iface = gr.Interface(
     fn=predict_depth,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
     title="MiDaS Depth Estimation",
-    description="Загрузите изображение и получите карту глубины."
+    description="Загрузите изображение или отправьте через API. Получите карту глубины."
 )
 
 if __name__ == "__main__":