Update gradio_app.py

gradio_app.py

@@ -57,25 +57,26 @@ if not hasattr(Image, 'Resampling'):
 
 
 def sam_init():
-    model = SamModel.from_pretrained("facebook/sam-vit-huge")
+    model = SamModel.from_pretrained("facebook/sam-vit-huge").to("cuda")
     processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
     return model, processor
 
 def sam_segment(sam_model, sam_processor, input_image, *bbox_coords):
     bbox = torch.tensor(bbox_coords, dtype=torch.float32)
-    bbox = box_tensor = bbox.unsqueeze(0).unsqueeze(0) # (1, 1, 4)
+    bbox =  bbox.unsqueeze(0)
     image = np.asarray(input_image)
 
     start_time = time.time()
 
     inputs = sam_processor(input_image, input_boxes=bbox, return_tensors="pt").to("cuda")
     outputs = sam_model(**inputs)
+    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] = outputs["masks"][-1].astype(np.uint8) * 255
+    out_image_bbox[:, :, 3] = masks[-1].astype(np.uint8) * 255
     torch.cuda.empty_cache()
     return Image.fromarray(out_image_bbox, mode='RGBA')