Update app.py

app.py

@@ -56,7 +56,9 @@ def predict(input, topk):
     t_image = img_resize.apply_image_torch(image)
     t_orig_size = t_image.shape[-2:]
     # pad to 1024x1024
+    pixel_mask = torch.ones(1, t_orig_size[0], t_orig_size[1], device=device)
     t_image = torch.nn.functional.pad(t_image, (0, 1024 - t_image.shape[-1], 0, 1024 - t_image.shape[-2]))
+    pixel_mask = torch.nn.functional.pad(pixel_mask, (0, 1024 - t_orig_size[1], 0, 1024 - t_orig_size[0]))
 
     # get box prompt
     valid_boxes = []
@@ -69,7 +71,7 @@ def predict(input, topk):
     t_boxes = np.array(valid_boxes)
     t_boxes = img_resize.apply_boxes(t_boxes, orig_size)
     box_torch = torch.as_tensor(t_boxes, dtype=torch.float, device=device)
-    batched_inputs = [{"image": t_image[0], "boxes": box_torch}]
+    batched_inputs = [{"image": t_image[0], "boxes": box_torch, "pixel_mask": pixel_mask}]
     with torch.no_grad():
         outputs = sam.infer(batched_inputs, multimask_output=False)
     # visualize and post on tensorboard
@@ -87,7 +89,7 @@ def predict(input, topk):
         pred_logits = outputs.logits[i].detach().cpu().numpy()
         top_ind = pred_logits[:, 0].argsort()[-topk:][::-1]
         pred_grasp = outputs.pred_boxes[i].detach().cpu().numpy()[top_ind]
-        coded_grasp = GraspCoder(1024, 1024, None, grasp_annos_reformat=pred_grasp)
+        coded_grasp = GraspCoder(t_orig_size[0], t_orig_size[1], None, grasp_annos_reformat=pred_grasp)
         _ = coded_grasp.decode()
         decoded_grasp = copy.deepcopy(coded_grasp.grasp_annos)
 
@@ -125,7 +127,4 @@ if __name__ == "__main__":
         btn.click(predict,
                   inputs=[prompter, top_k],
                   outputs=[image_output])
-    app.launch()
-
-
-
+    app.launch()