removing bg first

main.py

@@ -7,6 +7,7 @@ import cv2
 from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
 from PIL import Image
 import zipfile
+from transformers import pipeline
 
 app = Flask(__name__)
 CORS(app)
@@ -14,8 +15,8 @@ CORS(app)
 cudaOrNah = "cuda" if torch.cuda.is_available() else "cpu"
 print(cudaOrNah)
 
-# Global model setup
-# Adjusted due to memory constraints
+# Global model setup 
+# running out of memory adjusted
 # checkpoint = "sam_vit_h_4b8939.pth"
 # model_type = "vit_h"
 checkpoint = "sam_vit_l_0b3195.pth"
@@ -27,10 +28,12 @@ mask_generator = SamAutomaticMaskGenerator(
     min_mask_region_area=0.0015  # Adjust this value as needed
 )
 print('Setup SAM model')
+rembg_pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
+print('Setup rembg model')
 
 @app.route('/')
 def hello():
-    return {"hei": "Shredded to pieces"}
+    return {"hei": "Shredded to peices"}
 
 @app.route('/health', methods=['GET'])
 def health_check():
@@ -44,7 +47,7 @@ def get_masks():
         # Get the image file from the request
         if 'image' not in request.files:
             return jsonify({"error": "No image file provided"}), 400
-
+        
         image_file = request.files['image']
         if image_file.filename == '':
             return jsonify({"error": "No image file provided"}), 400
@@ -57,42 +60,38 @@ def get_masks():
 
         if image is None:
             raise ValueError("Image not found or unable to read.")
-
+        
         if cudaOrNah == "cuda":
             torch.cuda.empty_cache()
-
-        masks = mask_generator.generate(image)
+        
+        noBg = rembg_pipe(image)
+        masks = mask_generator.generate(noBg)
 
         if cudaOrNah == "cuda":
             torch.cuda.empty_cache()
 
-        # Sort masks by area in descending order
-        masks = sorted(masks, key=lambda x: x['area'], reverse=True)
-
-        # Initialize a cumulative mask to keep track of covered areas
-        cumulative_mask = np.zeros_like(masks[0]['segmentation'], dtype=bool)
-
-        # Process masks to remove overlaps
-        for mask in masks:
-            # Subtract areas already covered
-            mask['segmentation'] = np.logical_and(
-                mask['segmentation'], np.logical_not(cumulative_mask)
-            )
-            # Update the cumulative mask
-            cumulative_mask = np.logical_or(cumulative_mask, mask['segmentation'])
-            # Update the area
-            mask['area'] = mask['segmentation'].sum()
+        masks = sorted(masks, key=(lambda x: x['area']), reverse=True)
 
-        # Remove masks with zero area
-        masks = [mask for mask in masks if mask['area'] > 0]
-
-        # (Optional) Remove background masks if needed
         def is_background(segmentation):
             val = (segmentation[10, 10] or segmentation[-10, 10] or
                    segmentation[10, -10] or segmentation[-10, -10])
             return val
 
-        masks = [mask for mask in masks if not is_background(mask['segmentation'])]
+        # masks = [mask for mask in masks if not is_background(mask['segmentation'])]
+
+        for i in range(0, len(masks) - 1)[::-1]:
+            large_mask = masks[i]['segmentation']
+            for j in range(i+1, len(masks)):
+                not_small_mask = np.logical_not(masks[j]['segmentation'])
+                masks[i]['segmentation'] = np.logical_and(large_mask, not_small_mask)
+                masks[i]['area'] = masks[i]['segmentation'].sum()
+                large_mask = masks[i]['segmentation']
+
+        def sum_under_threshold(segmentation, threshold):
+            return segmentation.sum() / segmentation.size < 0.0015
+
+        masks = [mask for mask in masks if not sum_under_threshold(mask['segmentation'], 100)]
+        masks = sorted(masks, key=(lambda x: x['area']), reverse=True)
 
         # Create a zip file in memory
         zip_buffer = io.BytesIO()
@@ -106,7 +105,7 @@ def get_masks():
                 zip_file.writestr(f'mask_{idx+1}.png', mask_io.read())
 
         zip_buffer.seek(0)
-
+        
         return send_file(zip_buffer, mimetype='application/zip', as_attachment=True, download_name='masks.zip')
     except Exception as e:
         # Log the error message if needed
@@ -115,4 +114,4 @@ def get_masks():
         return jsonify({"error": "Error processing the image", "details": str(e)}), 400
 
 if __name__ == '__main__':
-    app.run(debug=True)
+    app.run(debug=True)

requirements.txt

@@ -12,4 +12,5 @@ torchvision
 matplotlib  # Required for image processing and mask visualization
 onnxruntime
 onnx
-pycocotools
+pycocotools
+transformers