Update app.py

app.py

@@ -1,142 +1,128 @@
 import torch
 import torchvision.transforms as transforms
-import gradio as gr
 import numpy as np
 import open3d as o3d
 from PIL import Image
 import cv2
-import uuid
-import gc  # Garbage collection
+import gc
 from transformers import DPTForDepthEstimation, DPTFeatureExtractor
+from torchvision.models.segmentation import deeplabv3_resnet50  # Optional for segmentation
 
-class RobustDepthTo3DConverter:
+class MultiModel3DReconstruction:
     def __init__(self):
-        # Explicitly use CPU
         self.device = torch.device("cpu")
         
         try:
-            # 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()
+            # Load DPT model for depth estimation
+            self.dpt_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(self.device)
+            self.dpt_feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
+            self.dpt_model.eval()
             
-            print("DPT model successfully initialized")
+            print("DPT model successfully loaded.")
+        except Exception as e:
+            print(f"Error loading DPT model: {e}")
+            self.dpt_model = None
         
+        try:
+            # Optional fallback: Load ZoeDepth model
+            self.zoe_model = torch.hub.load("isl-org/ZoeDepth", "ZoeD_N", pretrained=True).to(self.device)
+            self.zoe_model.eval()
+            
+            print("ZoeDepth model successfully loaded.")
         except Exception as e:
-            print(f"Critical model initialization error: {e}")
-            self.model = None
+            print(f"Error loading ZoeDepth model: {e}")
+            self.zoe_model = None
         
     def preprocess_image(self, input_image):
         """
-        Preprocess image using Hugging Face's feature extractor
+        Preprocess input image for models.
         """
-        # Ensure input is PIL Image
         if not isinstance(input_image, Image.Image):
             input_image = Image.fromarray(input_image)
         
-        # Convert to RGB if needed
-        if input_image.mode != 'RGB':
-            input_image = input_image.convert('RGB')
+        if input_image.mode != "RGB":
+            input_image = input_image.convert("RGB")
         
         return input_image
     
     def estimate_depth(self, input_image):
         """
-        Estimate depth using Hugging Face DPT model
+        Estimate depth using the best available model.
         """
-        if self.model is None:
-            raise ValueError("DPT model not properly initialized. Check model loading.")
+        input_image = self.preprocess_image(input_image)
         
         try:
-            # Preprocess the image
-            img = self.preprocess_image(input_image)
-            inputs = self.feature_extractor(images=img, return_tensors="pt").to(self.device)
-            
-            # Estimate depth
+            # Use DPT for depth estimation
+            inputs = self.dpt_feature_extractor(images=input_image, return_tensors="pt").to(self.device)
             with torch.no_grad():
-                outputs = self.model(**inputs)
+                outputs = self.dpt_model(**inputs)
                 depth = outputs.predicted_depth.squeeze().cpu().numpy()
-            
-            # Normalize depth
-            depth_normalized = cv2.normalize(
-                depth, None, 0, 255, 
-                norm_type=cv2.NORM_MINMAX, 
-                dtype=cv2.CV_8U
-            )
-            
-            # Manual memory cleanup
-            torch.cuda.empty_cache()
-            gc.collect()
-            
-            return depth_normalized
+                return depth
         
-        except Exception as e:
-            print(f"Depth estimation error: {e}")
+        except Exception as dpt_error:
+            print(f"DPT model error: {dpt_error}")
+        
+        try:
+            # Fallback: Use ZoeDepth for depth estimation
+            zoe_input = transforms.ToTensor()(input_image).unsqueeze(0).to(self.device)
+            with torch.no_grad():
+                depth = self.zoe_model.infer(zoe_input).squeeze().cpu().numpy()
+                return depth
+        
+        except Exception as zoe_error:
+            print(f"ZoeDepth fallback error: {zoe_error}")
+        
+        return None
+    
+    def refine_depth(self, depth_map):
+        """
+        Smooth and refine the depth map.
+        """
+        if depth_map is None:
             return None
+        
+        depth_map_normalized = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
+        refined = cv2.bilateralFilter(depth_map_normalized, d=9, sigmaColor=75, sigmaSpace=75)
+        return refined
     
     def create_point_cloud(self, image, depth_map):
         """
-        Create point cloud with reduced resolution
+        Generate a point cloud from the depth map.
         """
         if depth_map is None:
             return None
+        
+        img_array = np.array(image)
+        depth_map_resized = cv2.resize(depth_map, (img_array.shape[1], img_array.shape[0]), interpolation=cv2.INTER_LINEAR)
+        height, width = img_array.shape[:2]
 
-        try:
-            img_array = np.array(image)
-            depth_map_resized = cv2.resize(depth_map, (img_array.shape[1], img_array.shape[0]), interpolation=cv2.INTER_LINEAR)
-            height, width = img_array.shape[:2]
-
-            # Increase step size to reduce point cloud density
-            step = max(1, min(height, width) // 100)
-            points, colors = [], []
-
-            for y in range(0, height, step):
-                for x in range(0, width, step):
-                    z = depth_map_resized[y, x] / 255.0 * 3  # Reduced depth scale
-                    points.append([x, y, z])
-                    color = img_array[y, x][:3] / 255.0 if len(img_array[y, x]) >= 3 else [0.5, 0.5, 0.5]
-                    colors.append(color)
-
-            pcd = o3d.geometry.PointCloud()
-            pcd.points = o3d.utility.Vector3dVector(points)
-            pcd.colors = o3d.utility.Vector3dVector(colors)
-
-            return pcd
-
-        except Exception as e:
-            print(f"Point cloud creation error: {e}")
-            return None
+        step = max(1, min(height, width) // 100)  # Adjustable step size
+        points, colors = [], []
+        
+        for y in range(0, height, step):
+            for x in range(0, width, step):
+                z = depth_map_resized[y, x] / 255.0 * 3  # Adjust depth scaling
+                points.append([x, y, z])
+                color = img_array[y, x][:3] / 255.0 if len(img_array[y, x]) >= 3 else [0.5, 0.5, 0.5]
+                colors.append(color)
+        
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(points)
+        pcd.colors = o3d.utility.Vector3dVector(colors)
+        
+        return pcd
     
     def convert_to_mesh(self, point_cloud):
         """
-        More robust mesh conversion with reduced complexity
+        Convert a point cloud to a 3D mesh.
         """
         if point_cloud is None:
             return None
         
         try:
-            # Estimate normals with error handling
             point_cloud.estimate_normals()
-            point_cloud.orient_normals_consistent_tangent_plane(100)
-            
-            # More flexible mesh reconstruction with lower depth
-            try:
-                mesh, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
-                    point_cloud, depth=6  # Reduced depth for less memory usage
-                )
-            except RuntimeError:
-                # Fallback method if Poisson reconstruction fails
-                mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_ball_pivoting(
-                    point_cloud, o3d.utility.DoubleVector([0.02, 0.04, 0.08])
-                )
-            
-            # Simplify mesh to reduce memory and file size
+            mesh, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(point_cloud, depth=8)
             mesh = mesh.simplify_quadric_decimation(target_number_of_triangles=10000)
-            
-            # Color the mesh safely
-            if hasattr(point_cloud, 'colors'):
-                mesh.vertex_colors = point_cloud.colors
-            
             return mesh
         
         except Exception as e:
@@ -145,72 +131,31 @@ class RobustDepthTo3DConverter:
     
     def process_image(self, input_image):
         """
-        Enhanced full pipeline with comprehensive error handling and memory optimization
+        Complete pipeline: Depth estimation -> Point Cloud -> Mesh.
         """
-        # First, check if model is initialized
-        if self.model is None:
-            raise ValueError("DPT model initialization failed. Cannot process image.")
-        
         try:
-            # Preprocess and validate input
-            input_image = self.preprocess_image(input_image)
-            
-            # Estimate depth
             depth_map = self.estimate_depth(input_image)
             if depth_map is None:
-                raise ValueError("Depth estimation failed")
+                raise ValueError("Depth estimation failed.")
+            
+            refined_depth = self.refine_depth(depth_map)
+            point_cloud = self.create_point_cloud(input_image, refined_depth)
             
-            # Create point cloud
-            point_cloud = self.create_point_cloud(input_image, depth_map)
             if point_cloud is None:
-                raise ValueError("Point cloud creation failed")
+                raise ValueError("Point cloud generation failed.")
             
-            # Convert to mesh
             mesh = self.convert_to_mesh(point_cloud)
             if mesh is None:
-                raise ValueError("Mesh conversion failed")
-            
-            # Save mesh with unique filename
-            output_filename = f"/tmp/3d_model_{uuid.uuid4()}.obj"
-            o3d.io.write_triangle_mesh(output_filename, mesh)
+                raise ValueError("Mesh conversion failed.")
             
-            # Manual cleanup
-            del depth_map, point_cloud, mesh
-            gc.collect()
+            output_file = f"/tmp/3d_model_{uuid.uuid4()}.obj"
+            o3d.io.write_triangle_mesh(output_file, mesh)
             
-            return output_filename
+            return output_file
         
         except Exception as e:
-            print(f"Full pipeline error: {e}")
-            raise  # Re-raise the exception to be caught in the Gradio interface
-
-def create_huggingface_space():
-    # Initialize converter
-    converter = RobustDepthTo3DConverter()
-    
-    def convert_image(input_image):
-        try:
-            # Check model initialization before processing
-            if converter.model is None:
-                raise ValueError("DPT model failed to initialize. Cannot process image.")
-            
-            output_model = converter.process_image(input_image)
-            return output_model
-        except Exception as e:
-            print(f"Conversion error: {e}")
-            raise gr.Error(f"Conversion failed: {str(e)}")
-    
-    # Gradio Interface
-    iface = gr.Interface(
-        fn=convert_image,
-        inputs=gr.Image(type="pil", label="Input Image"),
-        outputs=gr.File(label="3D Model (OBJ)"),
-        title="Optimized AI Image to 3D Converter",
-        description="Convert images to 3D models with CPU optimization and reduced memory usage."
-    )
-    
-    return iface
+            print(f"Pipeline error: {e}")
+            return None
 
-# Launch the Gradio interface
-demo = create_huggingface_space()
-demo.launch(debug=True)
+# Instantiate and test the pipeline
+converter = MultiModel3DReconstruction()