l

main.py

@@ -14,9 +14,10 @@ CORS(app)
 cudaOrNah = "cuda" if torch.cuda.is_available() else "cpu"
 print(cudaOrNah)
 
-# Global model setup
-# Adjust the model type and checkpoint as needed
-# For example, using the "vit_l" model
+# Global model setup 
+# running out of memory adjusted
+# checkpoint = "sam_vit_h_4b8939.pth"
+# model_type = "vit_h"
 checkpoint = "sam_vit_l_0b3195.pth"
 model_type = "vit_l"
 sam = sam_model_registry[model_type](checkpoint=checkpoint)
@@ -29,7 +30,7 @@ print('Setup SAM model')
 
 @app.route('/')
 def hello():
-    return {"hei": "Shredded to pieces"}
+    return {"hei": "Shredded to peices"}
 
 @app.route('/health', methods=['GET'])
 def health_check():
@@ -39,45 +40,27 @@ 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
 
-        # **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)
+        raw_image = Image.open(image_file).convert("RGB")
+        # Convert the PIL Image to a NumPy array
         image_array = np.array(raw_image)
-
-        # 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
+        # Since OpenCV expects BGR, 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":
@@ -85,23 +68,28 @@ def get_masks():
 
         masks = sorted(masks, key=(lambda x: x['area']), reverse=True)
 
-        # 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)):
-                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 is_background(segmentation):
+            val = (segmentation[10, 10] or segmentation[-10, 10] or
+                   segmentation[10, -10] or segmentation[-10, -10])
+            return val
 
-        # Filter out very small masks based on area threshold
-        def sum_under_threshold(segmentation, threshold):
-            return segmentation.sum() / segmentation.size < threshold
+        masks = [mask for mask in masks if not is_background(mask['segmentation'])]
 
-        masks = [mask for mask in masks if not sum_under_threshold(mask['segmentation'], 0.0015)]
+        # 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 containing the mask images
+        # Create a zip file in memory
         zip_buffer = io.BytesIO()
         with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
             for idx, mask in enumerate(masks):
@@ -110,10 +98,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
@@ -122,4 +110,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)