dismantle

main.py

@@ -8,18 +8,12 @@ from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
 from PIL import Image
 import zipfile
 
-import os
-os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
-
-from transformers import pipeline
-
 app = Flask(__name__)
 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"
@@ -33,8 +27,6 @@ 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():
@@ -69,8 +61,7 @@ def get_masks():
         if cudaOrNah == "cuda":
             torch.cuda.empty_cache()
         
-        noBg = rembg_pipe(image)
-        masks = mask_generator.generate(noBg)
+        masks = mask_generator.generate(image)
 
         if cudaOrNah == "cuda":
             torch.cuda.empty_cache()
@@ -82,7 +73,7 @@ def get_masks():
                    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']