more debugging to address possible GPU OOM or timeout for 3D generation

app.py

@@ -70,95 +70,106 @@ def resize_image(image_path, max_size=1024):
             img.save(temp_file, format="PNG")
             return temp_file.name
 
-@spaces.GPU(duration=20)
+@spaces.GPU(duration=30)  # Increased duration to 30 seconds
 def generate_3d_model(depth, image_path, focallength_px, simplification_factor=0.8, smoothing_iterations=1, thin_threshold=0.01):
     """
     Generate a textured 3D mesh from the depth map and the original image.
     """
-    # Load the RGB image and convert to a NumPy array
-    image = np.array(Image.open(image_path))
-    
-    # Ensure depth is a NumPy array
-    if isinstance(depth, torch.Tensor):
-        depth = depth.cpu().numpy()
-    
-    # Resize depth to match image dimensions if necessary
-    if depth.shape != image.shape[:2]:
-        depth = cv2.resize(depth, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_LINEAR)
-    
-    height, width = depth.shape
-
-    print(f"3D model generation - Depth shape: {depth.shape}")
-    print(f"3D model generation - Image shape: {image.shape}")
-
-    # Compute camera intrinsic parameters
-    fx = fy = float(focallength_px)  # Ensure focallength_px is a float
-    cx, cy = width / 2, height / 2  # Principal point at the image center
-
-    # Create a grid of (u, v) pixel coordinates
-    u = np.arange(0, width)
-    v = np.arange(0, height)
-    uu, vv = np.meshgrid(u, v)
-
-    # Convert pixel coordinates to real-world 3D coordinates using the pinhole camera model
-    Z = depth.flatten()
-    X = ((uu.flatten() - cx) * Z) / fx
-    Y = ((vv.flatten() - cy) * Z) / fy
-
-    # Stack the coordinates to form vertices (X, Y, Z)
-    vertices = np.vstack((X, Y, Z)).T
-
-    # Normalize RGB colors to [0, 1] for vertex coloring
-    colors = image.reshape(-1, 3) / 255.0
-
-    # Generate faces by connecting adjacent vertices to form triangles
-    faces = []
-    for i in range(height - 1):
-        for j in range(width - 1):
-            idx = i * width + j
-            # Triangle 1
-            faces.append([idx, idx + width, idx + 1])
-            # Triangle 2
-            faces.append([idx + 1, idx + width, idx + width + 1])
-    faces = np.array(faces)
-
-    # Create the mesh using Trimesh with vertex colors
-    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, vertex_colors=colors)
-
-    # Mesh cleaning and improvement steps
-    print("Original mesh - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
-
-    # 1. Mesh simplification
-    target_faces = int(len(mesh.faces) * simplification_factor)
-    mesh = mesh.simplify_quadric_decimation(face_count=target_faces)
-    print("After simplification - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
-
-    # 2. Remove small disconnected components
-    components = mesh.split(only_watertight=False)
-    if len(components) > 1:
-        areas = np.array([c.area for c in components])
-        mesh = components[np.argmax(areas)]
-        print("After removing small components - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
-
-    # 3. Smooth the mesh
-    for _ in range(smoothing_iterations):
-        mesh = mesh.smoothed()
-    print("After smoothing - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
-
-    # 4. Remove thin features
-    mesh = remove_thin_features(mesh, thickness_threshold=thin_threshold)
-    print("After removing thin features - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
-
-    # Export the mesh to OBJ files with unique filenames
-    timestamp = int(time.time())
-    view_model_path = f'view_model_{timestamp}.obj'
-    download_model_path = f'download_model_{timestamp}.obj'
-    print("gonna export to view!")
-    mesh.export(view_model_path)
-    print("gonna export to download!")
-    mesh.export(download_model_path)
-    print("exported!")
-    return view_model_path, download_model_path
+    try:
+        print("Starting 3D model generation")
+        # Load the RGB image and convert to a NumPy array
+        image = np.array(Image.open(image_path))
+        
+        # Ensure depth is a NumPy array
+        if isinstance(depth, torch.Tensor):
+            depth = depth.cpu().numpy()
+        
+        # Resize depth to match image dimensions if necessary
+        if depth.shape != image.shape[:2]:
+            depth = cv2.resize(depth, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_LINEAR)
+        
+        height, width = depth.shape
+
+        print(f"3D model generation - Depth shape: {depth.shape}")
+        print(f"3D model generation - Image shape: {image.shape}")
+
+        # Compute camera intrinsic parameters
+        fx = fy = float(focallength_px)  # Ensure focallength_px is a float
+        cx, cy = width / 2, height / 2  # Principal point at the image center
+
+        # Create a grid of (u, v) pixel coordinates
+        u = np.arange(0, width)
+        v = np.arange(0, height)
+        uu, vv = np.meshgrid(u, v)
+
+        # Convert pixel coordinates to real-world 3D coordinates using the pinhole camera model
+        Z = depth.flatten()
+        X = ((uu.flatten() - cx) * Z) / fx
+        Y = ((vv.flatten() - cy) * Z) / fy
+
+        # Stack the coordinates to form vertices (X, Y, Z)
+        vertices = np.vstack((X, Y, Z)).T
+
+        # Normalize RGB colors to [0, 1] for vertex coloring
+        colors = image.reshape(-1, 3) / 255.0
+
+        print("Generating faces")
+        # Generate faces by connecting adjacent vertices to form triangles
+        faces = []
+        for i in range(height - 1):
+            for j in range(width - 1):
+                idx = i * width + j
+                # Triangle 1
+                faces.append([idx, idx + width, idx + 1])
+                # Triangle 2
+                faces.append([idx + 1, idx + width, idx + width + 1])
+        faces = np.array(faces)
+
+        print("Creating mesh")
+        # Create the mesh using Trimesh with vertex colors
+        mesh = trimesh.Trimesh(vertices=vertices, faces=faces, vertex_colors=colors)
+
+        # Mesh cleaning and improvement steps
+        print("Original mesh - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
+
+        print("Simplifying mesh")
+        # 1. Mesh simplification
+        target_faces = int(len(mesh.faces) * simplification_factor)
+        mesh = mesh.simplify_quadric_decimation(face_count=target_faces)
+        print("After simplification - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
+
+        print("Removing small components")
+        # 2. Remove small disconnected components
+        components = mesh.split(only_watertight=False)
+        if len(components) > 1:
+            areas = np.array([c.area for c in components])
+            mesh = components[np.argmax(areas)]
+            print("After removing small components - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
+
+        print("Smoothing mesh")
+        # 3. Smooth the mesh
+        for _ in range(smoothing_iterations):
+            mesh = mesh.smoothed()
+        print("After smoothing - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
+
+        print("Removing thin features")
+        # 4. Remove thin features
+        mesh = remove_thin_features(mesh, thickness_threshold=thin_threshold)
+        print("After removing thin features - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
+
+        # Export the mesh to OBJ files with unique filenames
+        timestamp = int(time.time())
+        view_model_path = f'view_model_{timestamp}.obj'
+        download_model_path = f'download_model_{timestamp}.obj'
+        print("Exporting to view")
+        mesh.export(view_model_path)
+        print("Exporting to download")
+        mesh.export(download_model_path)
+        print("Export completed")
+        return view_model_path, download_model_path
+    except Exception as e:
+        print(f"Error in generate_3d_model: {str(e)}")
+        raise
 
 def remove_thin_features(mesh, thickness_threshold=0.01):
     """