Changing ai model

app.py

@@ -7,6 +7,7 @@ from PIL import Image
 import cv2
 import uuid
 import gc  # Garbage collection
+from transformers import DPTForDepthEstimation, DPTFeatureExtractor
 
 class RobustDepthTo3DConverter:
     def __init__(self):
@@ -14,32 +15,20 @@ class RobustDepthTo3DConverter:
         self.device = torch.device("cpu")
         
         try:
-            # Use the smallest available MiDaS model for less memory usage
-            self.model = torch.hub.load("intel-isl/MiDaS", "MiDaS_small", pretrained=True, force_reload=False)
-            
-            # Move model to device and set to eval mode
-            self.model.to(self.device)
+            # Load Hugging Face DPT model and feature extractor
+            self.model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(self.device)
+            self.feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
             self.model.eval()
             
-            print("MiDaS model successfully initialized")
+            print("DPT model successfully initialized")
         
         except Exception as e:
             print(f"Critical model initialization error: {e}")
             self.model = None
         
-        # Create transformation pipeline with smaller input size
-        self.transform = transforms.Compose([
-            transforms.Resize((192, 192)),  # Reduced input size
-            transforms.ToTensor(),
-            transforms.Normalize(
-                mean=[0.485, 0.456, 0.406], 
-                std=[0.229, 0.224, 0.225]
-            )
-        ])
-    
     def preprocess_image(self, input_image):
         """
-        Standardize image input with more robust preprocessing
+        Preprocess image using Hugging Face's feature extractor
         """
         # Ensure input is PIL Image
         if not isinstance(input_image, Image.Image):
@@ -53,27 +42,20 @@ class RobustDepthTo3DConverter:
     
     def estimate_depth(self, input_image):
         """
-        More robust depth estimation with memory optimization
+        Estimate depth using Hugging Face DPT model
         """
         if self.model is None:
-            raise ValueError("MiDaS model not properly initialized. Check model loading.")
+            raise ValueError("DPT model not properly initialized. Check model loading.")
         
         try:
-            # Preprocess image
+            # Preprocess the image
             img = self.preprocess_image(input_image)
-            
-            # Convert to tensor
-            img_tensor = self.transform(img).unsqueeze(0).to(self.device)
+            inputs = self.feature_extractor(images=img, return_tensors="pt").to(self.device)
             
             # Estimate depth
             with torch.no_grad():
-                prediction = self.model(img_tensor)
-                
-                # Handle different model output formats
-                if isinstance(prediction, tuple):
-                    depth = prediction[0].squeeze().cpu().numpy()
-                else:
-                    depth = prediction.squeeze().cpu().numpy()
+                outputs = self.model(**inputs)
+                depth = outputs.predicted_depth.squeeze().cpu().numpy()
             
             # Normalize depth
             depth_normalized = cv2.normalize(
@@ -83,7 +65,6 @@ class RobustDepthTo3DConverter:
             )
             
             # Manual memory cleanup
-            del img_tensor
             torch.cuda.empty_cache()
             gc.collect()
             
@@ -168,7 +149,7 @@ class RobustDepthTo3DConverter:
         """
         # First, check if model is initialized
         if self.model is None:
-            raise ValueError("MiDaS model initialization failed. Cannot process image.")
+            raise ValueError("DPT model initialization failed. Cannot process image.")
         
         try:
             # Preprocess and validate input
@@ -211,7 +192,7 @@ def create_huggingface_space():
         try:
             # Check model initialization before processing
             if converter.model is None:
-                raise ValueError("MiDaS model failed to initialize. Cannot process image.")
+                raise ValueError("DPT model failed to initialize. Cannot process image.")
             
             output_model = converter.process_image(input_image)
             return output_model
@@ -232,4 +213,4 @@ def create_huggingface_space():
 
 # Launch the Gradio interface
 demo = create_huggingface_space()
-demo.launch(debug=True)
+demo.launch(debug=True)