Update app.py

app.py

@@ -17,6 +17,8 @@ import trimesh
 from transformers import pipeline
 from scipy.ndimage import gaussian_filter
 import open3d as o3d
+from rembg import remove
+import cv2
 
 # Force CPU usage
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -96,13 +98,28 @@ def process_with_timeout(function, args, timeout):
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
-# Image preprocessing for Depth-Anything (512x512, no background removal)
+# Image preprocessing: Remove background and resize
 def preprocess_image(image_path):
-    with Image.open(image_path) as img:
-        img = img.convert("RGB")
-        # Depth-Anything requires 512x512
-        img = img.resize((512, 512), Image.LANCZOS)
-        return img
+    try:
+        # Load image
+        with Image.open(image_path) as img:
+            # Remove background using rembg
+            img_no_bg = remove(img)
+            # Convert to RGB if it has an alpha channel
+            if img_no_bg.mode == 'RGBA':
+                img_no_bg = img_no_bg.convert('RGB')
+            # Resize to 512x512
+            img_no_bg = img_no_bg.resize((512, 512), Image.LANCZOS)
+            
+            # Optional: Use cv2 for additional segmentation (e.g., refine mask)
+            img_array = np.array(img_no_bg)
+            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+            _, mask = cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY)
+            img_array = cv2.bitwise_and(img_array, img_array, mask=mask)
+            
+            return Image.fromarray(img_array)
+    except Exception as e:
+        raise Exception(f"Error preprocessing image: {str(e)}")
 
 def load_model():
     global depth_pipeline, model_loaded, model_loading
@@ -182,20 +199,20 @@ def depth_to_point_cloud(depth_map, image, detail_level):
     h, w = depth_array.shape
     x, y = np.meshgrid(np.arange(w), np.arange(h))
     
-    # Simple camera intrinsics (assumed focal length)
+    # Camera intrinsics (assumed focal length)
     fx = fy = w * 0.5
     cx, cy = w / 2, h / 2
     
-    # Convert to 3D coordinates
+    # Convert to 3D coordinates (Z-up for Unity)
     z = depth_array
     x = (x - cx) * z / fx
-    y = (y - cy) * z / fy
+    y = -(y - cy) * z / fy  # Flip y-axis to correct orientation
     
     points = np.stack([x, y, z], axis=-1).reshape(-1, 3)
     colors = img_array.reshape(-1, 3) / 255.0
     
-    # Filter out invalid points (e.g., background)
-    mask = (z.reshape(-1) > 0.1) & (z.reshape(-1) < 0.9)
+    # Filter out invalid points (tighter range for foreground)
+    mask = (z.reshape(-1) > 0.2) & (z.reshape(-1) < 0.8)
     points = points[mask]
     colors = colors[mask]
     
@@ -223,6 +240,9 @@ def depth_to_point_cloud(depth_map, image, detail_level):
         vertex_colors=vertex_colors
     )
     
+    # Rotate mesh to correct orientation (180 degrees around X-axis)
+    trimesh_mesh.apply_transform(trimesh.transformations.rotation_matrix(np.pi, [1, 0, 0]))
+    
     return trimesh_mesh
 
 @app.route('/health', methods=['GET'])
@@ -297,7 +317,8 @@ def convert_image_to_3d():
     filepath = os.path.join(app.config['UPLOAD_FOLDER'], f"{job_id}_{filename}")
     file.save(filepath)
     
-    processing_jobs[job_id] = {
+    processing_jobs[job_id] =*.
+    {
         'status': 'processing',
         'progress': 0,
         'result_url': None,
@@ -484,7 +505,7 @@ def index():
             "output_format": "glb",
             "detail_level": "low, medium, or high - controls point cloud density"
         },
-        "description": "This API creates 3D models from 2D images using Depth-Anything depth estimation. Images should have transparent backgrounds."
+        "description": "This API creates 3D models from 2D images using Depth-Anything depth estimation. Images should have transparent backgrounds for best results."
     }), 200
 
 if __name__ == '__main__':