mask gen transparency ignore

main.py

@@ -14,10 +14,9 @@ CORS(app)
 cudaOrNah = "cuda" if torch.cuda.is_available() else "cpu"
 print(cudaOrNah)
 
-# Global model setup 
-# running out of memory adjusted
-# checkpoint = "sam_vit_h_4b8939.pth"
-# model_type = "vit_h"
+# Global model setup
+# Adjust the model type and checkpoint as needed
+# For example, using the "vit_l" model
 checkpoint = "sam_vit_l_0b3195.pth"
 model_type = "vit_l"
 sam = sam_model_registry[model_type](checkpoint=checkpoint)
@@ -30,7 +29,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():
@@ -40,27 +39,45 @@ def health_check():
 @app.route('/get-masks', methods=['POST'])
 def get_masks():
     try:
-        print('received image from frontend')
+        print('Received image from frontend')
         # 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
 
-        raw_image = Image.open(image_file).convert("RGB")
-        # Convert the PIL Image to a NumPy array
+        # **Modified Section Starts Here**
+
+        # Load the image with alpha channel to preserve transparency
+        raw_image = Image.open(image_file).convert("RGBA")
+        # Convert the PIL Image to a NumPy array (shape: H x W x 4)
         image_array = np.array(raw_image)
-        # Since OpenCV expects BGR, convert RGB to BGR
+
+        # Extract the alpha channel to create a transparency mask
+        alpha_channel = image_array[:, :, 3]
+        transparency_mask = alpha_channel > 0  # True where pixel is opaque
+
+        # Apply the transparency mask to the RGB channels
+        # Set transparent pixels to black (or any background color)
+        image_array[~transparency_mask, :3] = [0, 0, 0]
+
+        # Discard the alpha channel as it's no longer needed
+        image_array = image_array[:, :, :3]
+
+        # Since OpenCV expects BGR format, convert RGB to BGR
         image = image_array[:, :, ::-1]
 
+        # **Modified Section Ends Here**
+
         if image is None:
             raise ValueError("Image not found or unable to read.")
-        
+
         if cudaOrNah == "cuda":
             torch.cuda.empty_cache()
-        
+
+        # Generate masks using the preprocessed image
         masks = mask_generator.generate(image)
 
         if cudaOrNah == "cuda":
@@ -68,28 +85,23 @@ def get_masks():
 
         masks = sorted(masks, key=(lambda x: x['area']), reverse=True)
 
-        # 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'])]
-
-        for i in range(0, len(masks) - 1)[::-1]:
+        # Optional: Remove overlapping masks and small regions
+        for i in range(len(masks) - 1, -1, -1):
             large_mask = masks[i]['segmentation']
-            for j in range(i+1, len(masks)):
+            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']
 
+        # Filter out very small masks based on area threshold
         def sum_under_threshold(segmentation, threshold):
-            return segmentation.sum() / segmentation.size < 0.0015
+            return segmentation.sum() / segmentation.size < threshold
 
-        masks = [mask for mask in masks if not sum_under_threshold(mask['segmentation'], 100)]
+        masks = [mask for mask in masks if not sum_under_threshold(mask['segmentation'], 0.0015)]
         masks = sorted(masks, key=(lambda x: x['area']), reverse=True)
 
-        # Create a zip file in memory
+        # Create a zip file in memory containing the mask images
         zip_buffer = io.BytesIO()
         with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
             for idx, mask in enumerate(masks):
@@ -98,10 +110,10 @@ def get_masks():
                 mask_io = io.BytesIO()
                 mask_image.save(mask_io, format="PNG")
                 mask_io.seek(0)
-                zip_file.writestr(f'mask_{idx+1}.png', mask_io.read())
+                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
@@ -110,4 +122,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)