Update app.py

app.py

@@ -4,21 +4,20 @@ import cv2
 import numpy as np
 from torchvision import transforms
 from PIL import Image
-from transformers import DPTForDepthEstimation, DPTFeatureExtractor, MidasForDepthEstimation, MidasImageProcessor
 
-# Load depth estimation model (MiDaS v3 for better accuracy)
-model_name = "Intel/midas-v3"  # Upgraded model
-processor = MidasImageProcessor.from_pretrained(model_name)
-depth_model = MidasForDepthEstimation.from_pretrained(model_name)
-depth_model.eval()
+# Load MiDaS depth estimation model from torch.hub
+midas_model = torch.hub.load("intel-isl/MiDaS", "MiDaS_small")
+midas_model.eval()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+midas_model.to(device)
+midas_transform = torch.hub.load("intel-isl/MiDaS", "transforms").default_transform
 
 def estimate_depth(image):
-    """Estimate depth map from image using MiDaS v3."""
+    """Estimate depth map using MiDaS v3."""
     image = image.convert("RGB")
-    inputs = processor(images=image, return_tensors="pt")
+    input_tensor = midas_transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
-        outputs = depth_model(**inputs)
-        depth = outputs.predicted_depth.squeeze().cpu().numpy()
+        depth = midas_model(input_tensor).squeeze().cpu().numpy()
     depth = cv2.resize(depth, (image.width, image.height))
     depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255
     return depth.astype(np.uint8)
@@ -82,4 +81,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch(share=True)  
+    iface.launch(share=True)