Update app.py

app.py

@@ -16,7 +16,6 @@ from flask_cors import CORS
 import functools
 import numpy as np
 import trimesh
-from scipy.spatial import Delaunay
 
 app = Flask(__name__)
 CORS(app)  # Enable CORS for all routes
@@ -185,23 +184,89 @@ class NeuS2Model:
         
         sdf_volume = torch.cat(sdf_values, dim=1).reshape(resolution, resolution, resolution).cpu().numpy()
         
-        # Extract mesh using marching cubes
-        vertices, faces = self._marching_cubes(sdf_volume, threshold)
+        # Extract mesh - alternative to marching cubes since we don't have scipy
+        vertices, faces = self._simple_mesh_extraction(sdf_volume, threshold)
         
         # Create a mesh object with vertices and faces
         mesh = type('Mesh', (), {'verts': vertices, 'faces': faces})
         
         return [mesh]  # Returning in list format to match ShapE's output format
     
-    def _marching_cubes(self, sdf_volume, threshold=0.0):
-        # Simple implementation using surface points and Delaunay triangulation
-        # For production, you'd want to use proper marching cubes from scikit-image
+    def _simple_mesh_extraction(self, sdf_volume, threshold=0.0):
+        """Simple mesh extraction without scipy dependency"""
+        resolution = sdf_volume.shape[0]
         
-        # Find points near the surface
-        x, y, z = np.where(np.abs(sdf_volume) < 0.1)
+        # Find surface points (approximate)
+        surface_points = []
+        surface_normals = []
         
-        if len(x) < 4:  # Need at least 4 points for Delaunay
-            # Create a simple cube if not enough points
+        # Sample points on three orthogonal grids
+        for axis in range(3):
+            for i in range(resolution):
+                for j in range(resolution):
+                    # Create a line along the current axis
+                    line = np.zeros((resolution, 3), dtype=int)
+                    for k in range(resolution):
+                        if axis == 0:
+                            line[k] = [k, i, j]
+                        elif axis == 1:
+                            line[k] = [i, k, j]
+                        else:
+                            line[k] = [i, j, k]
+                    
+                    # Get SDF values along this line
+                    sdf_line = np.array([sdf_volume[tuple(idx)] for idx in line])
+                    
+                    # Find zero crossings
+                    signs = np.sign(sdf_line)
+                    zero_crossings = np.where(np.diff(signs) != 0)[0]
+                    
+                    for idx in zero_crossings:
+                        # Linear interpolation to find more accurate zero crossing
+                        t = sdf_line[idx] / (sdf_line[idx] - sdf_line[idx + 1])
+                        point = line[idx] * (1 - t) + line[idx + 1] * t
+                        
+                        # Normalize to [-1, 1] range
+                        normalized_point = 2 * point / resolution - 1
+                        surface_points.append(normalized_point)
+                        
+                        # Compute normal (gradient of SDF)
+                        normal = np.zeros(3)
+                        idx_3d = tuple(np.round(point).astype(int).clip(0, resolution - 1))
+                        
+                        # Compute gradient using central differences where possible
+                        for d in range(3):
+                            if 0 < idx_3d[d] < resolution - 1:
+                                idx_plus = list(idx_3d)
+                                idx_minus = list(idx_3d)
+                                idx_plus[d] += 1
+                                idx_minus[d] -= 1
+                                normal[d] = (sdf_volume[tuple(idx_plus)] - sdf_volume[tuple(idx_minus)]) / 2
+                            else:
+                                # Forward or backward difference at boundaries
+                                idx_curr = list(idx_3d)
+                                if idx_3d[d] == 0:
+                                    idx_other = list(idx_3d)
+                                    idx_other[d] = 1
+                                    normal[d] = sdf_volume[tuple(idx_other)] - sdf_volume[tuple(idx_curr)]
+                                else:
+                                    idx_other = list(idx_3d)
+                                    idx_other[d] = resolution - 2
+                                    normal[d] = sdf_volume[tuple(idx_curr)] - sdf_volume[tuple(idx_other)]
+                        
+                        if np.linalg.norm(normal) > 0:
+                            normal = normal / np.linalg.norm(normal)
+                            surface_normals.append(normal)
+        
+        # Limit the number of points to avoid OOM
+        max_points = 5000
+        if len(surface_points) > max_points:
+            indices = np.random.choice(len(surface_points), max_points, replace=False)
+            surface_points = [surface_points[i] for i in indices]
+            surface_normals = [surface_normals[i] for i in indices]
+        
+        if len(surface_points) < 4:
+            # Not enough points found, create a simple cube
             vertices = np.array([
                 [-0.5, -0.5, -0.5],
                 [0.5, -0.5, -0.5],
@@ -222,27 +287,104 @@ class NeuS2Model:
             ])
             return vertices, faces
         
-        # Convert indices to 3D coordinates in [-1, 1] range
-        res = sdf_volume.shape[0]
-        points = np.stack([
-            2 * x / res - 1,
-            2 * y / res - 1,
-            2 * z / res - 1
-        ], axis=1)
+        # Convert points to numpy array
+        points = np.array(surface_points)
         
-        # Limit to a reasonable number of points for Delaunay
-        max_points = 1000
-        if len(points) > max_points:
-            indices = np.random.choice(len(points), max_points, replace=False)
-            points = points[indices]
+        # Create a simple mesh using ball-pivoting like algorithm
+        vertices = points
         
-        try:
-            # Create triangular mesh using Delaunay
-            tri = Delaunay(points)
-            return points, tri.simplices
-        except Exception:
-            # Fallback to simple shape if Delaunay fails
-            return np.array([[-1, -1, -1], [1, -1, -1], [1, 1, -1], [-1, 1, -1]]), np.array([[0, 1, 2], [0, 2, 3]])
+        # For simplicity, create faces from nearest neighbors
+        # This is a very simple approach, not as good as Delaunay but doesn't require scipy
+        faces = []
+        
+        # Use a simple approach to creating faces
+        # We'll use a greedy algorithm to connect nearby points
+        n_points = len(vertices)
+        
+        # Create a simple connectivity graph
+        # For each point, find the N closest points
+        n_neighbors = min(12, n_points - 1)
+        adjacency = [[] for _ in range(n_points)]
+        
+        # Compute all pairwise distances (this is O(n²) but should be ok for small point clouds)
+        for i in range(n_points):
+            distances = []
+            for j in range(n_points):
+                if i != j:
+                    dist = np.linalg.norm(vertices[i] - vertices[j])
+                    distances.append((dist, j))
+            distances.sort()
+            for k in range(min(n_neighbors, len(distances))):
+                adjacency[i].append(distances[k][1])
+        
+        # Create triangles from the adjacency list
+        added_edges = set()
+        for i in range(n_points):
+            for j in adjacency[i]:
+                if j > i:  # Avoid duplicates
+                    edge_ij = (i, j)
+                    added_edges.add(edge_ij)
+                    
+                    # Find common neighbors between i and j to form triangles
+                    common_neighbors = set(adjacency[i]) & set(adjacency[j])
+                    for k in common_neighbors:
+                        if k != i and k != j:
+                            edge_ik = (i, k) if i < k else (k, i)
+                            edge_jk = (j, k) if j < k else (k, j)
+                            
+                            # Check if the other two edges exist
+                            if edge_ik in added_edges and edge_jk in added_edges:
+                                # Ensure consistent winding
+                                normal = np.cross(vertices[j] - vertices[i], vertices[k] - vertices[i])
+                                center_normal = np.mean(surface_normals, axis=0)
+                                
+                                if np.dot(normal, center_normal) > 0:
+                                    faces.append([i, j, k])
+                                else:
+                                    faces.append([i, k, j])
+        
+        # If we couldn't create enough faces, create a simple shape
+        if len(faces) < 4:
+            # Create a convex hull-like shape
+            center = np.mean(vertices, axis=0)
+            n_points = len(vertices)
+            
+            # Sort points by distance from center
+            dists = np.linalg.norm(vertices - center, axis=1)
+            sorted_indices = np.argsort(dists)
+            
+            # Create a star-like structure connecting to center
+            center_idx = sorted_indices[0]
+            faces = []
+            
+            for i in range(1, min(n_points, 10)):
+                if i + 1 < n_points:
+                    faces.append([center_idx, sorted_indices[i], sorted_indices[i + 1]])
+                else:
+                    faces.append([center_idx, sorted_indices[i], sorted_indices[1]])
+            
+            # If still not enough, create a simple cube
+            if len(faces) < 4:
+                vertices = np.array([
+                    [-0.5, -0.5, -0.5],
+                    [0.5, -0.5, -0.5],
+                    [0.5, 0.5, -0.5],
+                    [-0.5, 0.5, -0.5],
+                    [-0.5, -0.5, 0.5],
+                    [0.5, -0.5, 0.5],
+                    [0.5, 0.5, 0.5],
+                    [-0.5, 0.5, 0.5]
+                ])
+                faces = np.array([
+                    [0, 1, 2], [0, 2, 3],  # Bottom face
+                    [4, 5, 6], [4, 6, 7],  # Top face
+                    [0, 1, 5], [0, 5, 4],  # Front face
+                    [2, 3, 7], [2, 7, 6],  # Back face
+                    [0, 3, 7], [0, 7, 4],  # Left face
+                    [1, 2, 6], [1, 6, 5]   # Right face
+                ])
+        
+        return np.array(vertices), np.array(faces)
 
 def load_model():
     global neus_model, model_loaded, model_loading
@@ -347,6 +489,265 @@ def export_to_obj(mesh, obj_path):
     
     return obj_path, mtl_path
 
+@app.route('/convert', methods=['POST'])
+def convert_image_to_3d():
+    # Check if image is in the request
+    if 'image' not in request.files:
+        return jsonify({"error": "No image provided"}), 400
+    
+    file = request.files['image']
+    if file.filename == '':
+        return jsonify({"error": "No image selected"}), 400
+    
+    if not allowed_file(file.filename):
+        return jsonify({"error": f"File type not allowed. Supported types: {', '.join(ALLOWED_EXTENSIONS)}"}), 400
+    
+    # Get optional parameters with defaults
+    try:
+        guidance_scale = float(request.form.get('guidance_scale', 3.0))
+        num_inference_steps = min(int(request.form.get('num_inference_steps', 32)), MAX_INFERENCE_STEPS)
+        output_format = request.form.get('output_format', 'obj').lower()
+    except ValueError:
+        return jsonify({"error": "Invalid parameter values"}), 400
+    
+    # Validate parameters
+    if guidance_scale < 1.0 or guidance_scale > 5.0:
+        return jsonify({"error": "Guidance scale must be between 1.0 and 5.0"}), 400
+    
+    if num_inference_steps < 16 or num_inference_steps > MAX_INFERENCE_STEPS:
+        num_inference_steps = min(num_inference_steps, MAX_INFERENCE_STEPS)
+    
+    # Validate output format
+    if output_format not in ['obj', 'glb']:
+        return jsonify({"error": "Unsupported output format. Use 'obj' or 'glb'"}), 400
+    
+    # Create a job ID
+    job_id = str(uuid.uuid4())
+    output_dir = os.path.join(RESULTS_FOLDER, job_id)
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Save the uploaded file
+    filename = secure_filename(file.filename)
+    filepath = os.path.join(app.config['UPLOAD_FOLDER'], f"{job_id}_{filename}")
+    file.save(filepath)
+    
+    # Initialize job tracking
+    processing_jobs[job_id] = {
+        'status': 'processing',
+        'progress': 0,
+        'result_url': None,
+        'preview_url': None,
+        'error': None,
+        'output_format': output_format,
+        'created_at': time.time()
+    }
+    
+    # Start processing in a separate thread
+    def process_image():
+        thread = threading.current_thread()
+        processing_jobs[job_id]['thread_alive'] = lambda: thread.is_alive()
+        
+        try:
+            # Preprocess image (resize if needed)
+            processing_jobs[job_id]['progress'] = 5
+            image_tensor = preprocess_image(filepath)
+            processing_jobs[job_id]['progress'] = 10
+            
+            # Load model
+            try:
+                model = load_model()
+                processing_jobs[job_id]['progress'] = 30
+            except Exception as e:
+                processing_jobs[job_id]['status'] = 'error'
+                processing_jobs[job_id]['error'] = f"Error loading model: {str(e)}"
+                return
+            
+            # Process image with thread-safe timeout
+            try:
+                def generate_mesh():
+                    return model.generate_mesh(
+                        image_tensor, 
+                        resolution=min(32 + num_inference_steps // 2, 64),  # Adjust resolution based on steps
+                        threshold=0.0,
+                        num_steps=num_inference_steps
+                    )
+                
+                images, error = process_with_timeout(generate_mesh, [], TIMEOUT_SECONDS)
+                
+                if error:
+                    if isinstance(error, TimeoutError):
+                        processing_jobs[job_id]['status'] = 'error'
+                        processing_jobs[job_id]['error'] = f"Processing timed out after {TIMEOUT_SECONDS} seconds"
+                        return
+                    else:
+                        raise error
+                        
+                processing_jobs[job_id]['progress'] = 80
+            except Exception as e:
+                error_details = traceback.format_exc()
+                processing_jobs[job_id]['status'] = 'error'
+                processing_jobs[job_id]['error'] = f"Error during processing: {str(e)}"
+                print(f"Error processing job {job_id}: {str(e)}")
+                print(error_details)
+                return
+            
+            # Export based on requested format
+            try:
+                if output_format == 'obj':
+                    obj_path = os.path.join(output_dir, "model.obj")
+                    obj_path, mtl_path = export_to_obj(images[0], obj_path)
+                    
+                    # Create a zip file with OBJ and MTL
+                    zip_path = os.path.join(output_dir, "model.zip")
+                    with zipfile.ZipFile(zip_path, 'w') as zipf:
+                        zipf.write(obj_path, arcname="model.obj")
+                        if os.path.exists(mtl_path):
+                            zipf.write(mtl_path, arcname="model.mtl")
+                    
+                    processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
+                    processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
+                    
+                elif output_format == 'glb':
+                    # Convert to trimesh format
+                    vertices = images[0].verts
+                    faces = images[0].faces
+                    
+                    # Create a trimesh object
+                    trimesh_obj = trimesh.Trimesh(vertices=vertices, faces=faces)
+                    
+                    # Export as GLB
+                    glb_path = os.path.join(output_dir, "model.glb")
+                    trimesh_obj.export(glb_path)
+                    
+                    processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
+                    processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
+                
+                # Update job status
+                processing_jobs[job_id]['status'] = 'completed'
+                processing_jobs[job_id]['progress'] = 100
+                print(f"Job {job_id} completed successfully")
+            except Exception as e:
+                error_details = traceback.format_exc()
+                processing_jobs[job_id]['status'] = 'error'
+                processing_jobs[job_id]['error'] = f"Error exporting model: {str(e)}"
+                print(f"Error exporting model for job {job_id}: {str(e)}")
+                print(error_details)
+            
+            # Clean up temporary file
+            if os.path.exists(filepath):
+                os.remove(filepath)
+            
+            # Force garbage collection to free memory
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            
+        except Exception as e:
+            # Handle errors
+            error_details = traceback.format_exc()
+            processing_jobs[job_id]['status'] = 'error'
+            processing_jobs[job_id]['error'] = f"{str(e)}\n{error_details}"
+            print(f"Error processing job {job_id}: {str(e)}")
+            print(error_details)
+            
+            # Clean up on error
+            if os.path.exists(filepath):
+                os.remove(filepath)
+    
+    # Start processing thread
+    processing_thread = threading.Thread(target=process_image)
+    processing_thread.daemon = True
+    processing_thread.start()
+    
+    # Return job ID immediately
+    return jsonify({"job_id": job_id}), 202  # 202 Accepted
+
+@app.route('/download/<job_id>', methods=['GET'])
+def download_model(job_id):
+    if job_id not in processing_jobs or processing_jobs[job_id]['status'] != 'completed':
+        return jsonify({"error": "Model not found or processing not complete"}), 404
+    
+    # Get the output directory for this job
+    output_dir = os.path.join(RESULTS_FOLDER, job_id)
+    
+    # Determine file format from the job data
+    output_format = processing_jobs[job_id].get('output_format', 'obj')
+    
+    if output_format == 'obj':
+        zip_path = os.path.join(output_dir, "model.zip")
+        if os.path.exists(zip_path):
+            return send_file(zip_path, as_attachment=True, download_name="model.zip")
+    else:  # glb
+        glb_path = os.path.join(output_dir, "model.glb")
+        if os.path.exists(glb_path):
+            return send_file(glb_path, as_attachment=True, download_name="model.glb")
+    
+    return jsonify({"error": "File not found"}), 404
+@app.route('/progress/<job_id>', methods=['GET'])
+def progress(job_id):
+    def generate():
+        if job_id not in processing_jobs:
+            yield f"data: {json.dumps({'error': 'Job not found'})}\n\n"
+            return
+            
+        job = processing_jobs[job_id]
+        
+        # Send initial progress
+        yield f"data: {json.dumps({'status': 'processing', 'progress': job['progress']})}\n\n"
+        
+        # Wait for job to complete or update
+        last_progress = job['progress']
+        check_count = 0
+        while job['status'] == 'processing':
+            if job['progress'] != last_progress:
+                yield f"data: {json.dumps({'status': 'processing', 'progress': job['progress']})}\n\n"
+                last_progress = job['progress']
+            
+            time.sleep(0.5)
+            check_count += 1
+            
+            # If client hasn't received updates for a while, check if job is still running
+            if check_count > 60:  # 30 seconds with no updates
+                if 'thread_alive' in job and not job['thread_alive']():
+                    job['status'] = 'error'
+                    job['error'] = 'Processing thread died unexpectedly'
+                    break
+                check_count = 0
+        
+        # Send final status
+        if job['status'] == 'completed':
+            yield f"data: {json.dumps({'status': 'completed', 'progress': 100, 'result_url': job['result_url'], 'preview_url': job['preview_url']})}\n\n"
+        else:
+            yield f"data: {json.dumps({'status': 'error', 'error': job['error']})}\n\n"
+    
+    return Response(stream_with_context(generate()), mimetype='text/event-stream')
+
+def export_to_obj(mesh, obj_path):
+    """Export mesh to OBJ file format"""
+    vertices = mesh.verts
+    faces = mesh.faces
+    
+    with open(obj_path, 'w') as f:
+        # Write vertices
+        for v in vertices:
+            f.write(f"v {v[0]} {v[1]} {v[2]}\n")
+        
+        # Write faces (OBJ uses 1-indexed vertices)
+        for face in faces:
+            f.write(f"f {face[0]+1} {face[1]+1} {face[2]+1}\n")
+    
+    # Create a simple MTL file
+    mtl_path = obj_path.replace('.obj', '.mtl')
+    with open(mtl_path, 'w') as f:
+        f.write("newmtl material0\n")
+        f.write("Ka 1.0 1.0 1.0\n")  # ambient color
+        f.write("Kd 0.8 0.8 0.8\n")  # diffuse color
+        f.write("Ks 0.0 0.0 0.0\n")  # specular color
+        f.write("Ns 0.0\n")          # specular exponent
+        f.write("illum 2\n")         # illumination model
+    
+    return obj_path, mtl_path
+
 @app.route('/convert', methods=['POST'])
 def convert_image_to_3d():
     # Check if image is in the request