Update app.py

app.py

@@ -30,16 +30,23 @@ sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 def predict_masks_and_scores(model, processor, raw_image, input_points=None, input_boxes=None):
     if input_boxes is not None:
         input_boxes = [input_boxes]
-    inputs = processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt").to(model.device)
+
+    inputs = processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt")
+    original_sizes = inputs["original_sizes"]
+    reshaped_sizes = inputs["reshaped_input_sizes"]
+
+    inputs = inputs.to(model.device)
+
     with torch.no_grad():
         outputs = model(**inputs)
 
     masks = processor.image_processor.post_process_masks(
-        outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
+        outputs.pred_masks.cpu(), original_sizes, reshaped_sizes
     )
     scores = outputs.iou_scores
     return masks, scores
 
+
 def process_inputs(prompts):
     raw_entries = prompts["points"]