mask test 2

main.py

@@ -14,8 +14,8 @@ CORS(app)
 cudaOrNah = "cuda" if torch.cuda.is_available() else "cpu"
 print(cudaOrNah)
 
-# Global model setup 
-# running out of memory adjusted
+# Global model setup
+# Adjusted due to memory constraints
 # checkpoint = "sam_vit_h_4b8939.pth"
 # model_type = "vit_h"
 checkpoint = "sam_vit_l_0b3195.pth"
@@ -30,7 +30,7 @@ print('Setup SAM model')
 
 @app.route('/')
 def hello():
-    return {"hei": "Shredded to peices"}
+    return {"hei": "Shredded to pieces"}
 
 @app.route('/health', methods=['GET'])
 def health_check():
@@ -44,7 +44,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,17 +57,36 @@ 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)
 
         if cudaOrNah == "cuda":
             torch.cuda.empty_cache()
 
-        masks = sorted(masks, key=(lambda x: x['area']), reverse=True)
+        # 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()
 
+        # 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])
@@ -75,20 +94,6 @@ def get_masks():
 
         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()
         with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
@@ -101,7 +106,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