sam_segment

gradio_app.py

@@ -57,8 +57,8 @@ if not hasattr(Image, 'Resampling'):
 
 
 def sam_init():
-    model = SamModel.from_pretrained("facebook/sam-vit-huge").to("cuda")
-    processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
+    model = SamModel.from_pretrained("facebook/sam-vit-large").to("cuda")
+    processor = SamProcessor.from_pretrained("facebook/sam-vit-large")
     return model, processor
 
 def sam_segment(sam_model, sam_processor, input_image, *bbox_coords):
@@ -68,15 +68,17 @@ def sam_segment(sam_model, sam_processor, input_image, *bbox_coords):
 
     start_time = time.time()
 
-    inputs = sam_processor(input_image, input_boxes=bbox, return_tensors="pt").to("cuda")
-    outputs = sam_model(**inputs)
+    inputs = sam_processor(input_image.convert('RGB'), input_boxes=bbox, return_tensors="pt", do_resize=False).to("cuda")
+    
+    outputs = sam_model(**inputs, multimask_output=False)
     masks = sam_processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
 
     print(f"SAM Time: {time.time() - start_time:.3f}s")
     out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
     out_image[:, :, :3] = image
     out_image_bbox = out_image.copy()
-    out_image_bbox[:, :, 3] = masks[-1].astype(np.uint8) * 255
+
+    out_image_bbox[:, :, 3] = masks[-1].cpu().detach().numpy().astype(np.uint8) * 255
     torch.cuda.empty_cache()
     return Image.fromarray(out_image_bbox, mode='RGBA')