Update app.py

app.py

@@ -1,3 +1,4 @@
+
 import os
 import torch
 import time
@@ -17,7 +18,6 @@ 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
@@ -98,26 +98,55 @@ 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: Remove background and resize
+# Image preprocessing: Remove background using cv2
 def preprocess_image(image_path):
     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')
+            # Convert to RGB or handle transparency
+            if img.mode == 'RGBA':
+                # Use alpha channel as initial mask
+                img_array = np.array(img)
+                alpha = img_array[:, :, 3]
+                img_rgb = img_array[:, :, :3]
+            else:
+                img_rgb = np.array(img.convert('RGB'))
+                alpha = None
+            
             # Resize to 512x512
-            img_no_bg = img_no_bg.resize((512, 512), Image.LANCZOS)
+            img_rgb = cv2.resize(img_rgb, (512, 512), interpolation=cv2.INTER_LANCZOS4)
+            
+            # Convert to grayscale
+            gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
+            
+            # Adaptive thresholding for initial mask
+            thresh = cv2.adaptiveThreshold(
+                gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2
+            )
+            
+            # If alpha channel exists, combine with threshold
+            if alpha is not None:
+                alpha_resized = cv2.resize(alpha, (512, 512), interpolation=cv2.INTER_LANCZOS4)
+                thresh = cv2.bitwise_and(thresh, alpha_resized)
+            
+            # Refine with GrabCut
+            mask = np.zeros((512, 512), np.uint8)
+            mask[thresh == 255] = cv2.GC_PR_FGD  # Probable foreground
+            mask[thresh == 0] = cv2.GC_PR_BGD   # Probable background
+            
+            bgdModel = np.zeros((1, 65), np.float64)
+            fgdModel = np.zeros((1, 65), np.float64)
+            
+            rect = (10, 10, 492, 492)  # ROI for GrabCut
+            cv2.grabCut(img_rgb, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_MASK)
+            
+            # Create final mask (foreground = 1, background = 0)
+            mask2 = np.where((mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD), 255, 0).astype('uint8')
             
-            # Use cv2 for additional segmentation
-            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)
+            # Apply mask to image
+            img_foreground = cv2.bitwise_and(img_rgb, img_rgb, mask=mask2)
             
-            return Image.fromarray(img_array)
+            return Image.fromarray(img_foreground)
     except Exception as e:
         raise Exception(f"Error preprocessing image: {str(e)}")
 
@@ -516,4 +545,4 @@ def index():
 if __name__ == '__main__':
     cleanup_old_jobs()
     port = int(os.environ.get('PORT', 7860))
-    app.run(host='0.0.0.0', port=port)
+    app.run(host='0.0.0.0', port=port)