Using bounding box prompts

app.py

@@ -21,7 +21,19 @@ sidewalk_model.load_state_dict(checkpoint["model"])
 device = "cpu"
 sidewalk_model.to(device)
 # print('Status Update: Using GPU.')
+print('Status Update: FindMySidewalk Ready for inference ...')
 
+# Generate bounding box prompt for SAM
+def get_bounding_box(W = 256, H = 256, x_min = 0, y_min = 0, x_max = 256, y_max = 256):
+  # add perturbation if inputted bounding box coordinates
+  x_min = max(0, x_min - np.random.randint(0, 20))
+  x_max = min(W, x_max + np.random.randint(0, 20))
+  y_min = max(0, y_min - np.random.randint(0, 20))
+  y_max = min(H, y_max + np.random.randint(0, 20))
+
+  bbox = [x_min, y_min, x_max, y_max]
+
+  return bbox
 
 def segment_sidewalk(image):
     test_image = Image.fromarray(image).convert("RGB")
@@ -29,17 +41,19 @@ def segment_sidewalk(image):
     # Keep a copy of original image for display
     original_image = test_image.copy()
 
-    # Create grid of points for prompting
-    array_size = 256
-    grid_size = 7
-    x = np.linspace(0, array_size - 1, grid_size)
-    y = np.linspace(0, array_size - 1, grid_size)
-    xv, yv = np.meshgrid(x, y)
-    input_points = [[[int(x), int(y)] for x, y in zip(x_row, y_row)] for x_row, y_row in zip(xv.tolist(), yv.tolist())]
-    input_points = torch.tensor(input_points).view(1, 1, grid_size * grid_size, 2)
+    # # Create grid of points for prompting
+    # array_size = 256
+    # grid_size = 7
+    # x = np.linspace(0, array_size - 1, grid_size)
+    # y = np.linspace(0, array_size - 1, grid_size)
+    # xv, yv = np.meshgrid(x, y)
+    # input_points = [[[int(x), int(y)] for x, y in zip(x_row, y_row)] for x_row, y_row in zip(xv.tolist(), yv.tolist())]
+    # input_points = torch.tensor(input_points).view(1, 1, grid_size * grid_size, 2)
 
+    # obtain bounding box prompt over entire image
+    prompt = get_bounding_box(test_image.size[0], test_image.size[1], 0, 0, test_image.size[0], test_image.size[1])
     # prepare image for the model
-    inputs = processor(test_image, input_points=input_points, return_tensors="pt")
+    inputs = processor(test_image, input_boxes=[[prompt]], return_tensors="pt")
     # Convert dtype to float32 as the MPS framework doesn't support float64
     inputs = {k: v.to(torch.float32).to(device) for k, v in inputs.items()}
     sidewalk_model.eval()