SAM Model: Box and Point Input

[cm107]

In order to make it easier to generate the polygons that I desire, I would like to be able to specify a box input in addition to the point input.

It's pretty straightforward what needs to be done with the model's input. Refer to the Example box and point input section of onnx_model_example.ipynb.
For now I am using a modified version of _compute_polygon_from_points in order to accomplish this:

See Code

def _compute_polygon_from_points(
    image_size, decoder_session, image, image_embedding, points, point_labels
):
    points: list[list[float]] = points
    point_labels: list[int] = point_labels
    if len(points) < 3:
        input_point = np.array(points, dtype=np.float32)
        input_label = np.array(point_labels, dtype=np.int32)

        onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[
            None, :, :
        ]
        onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[
            None, :
        ].astype(np.float32)
    else:
        box_points = points[:2]
        points = points[2:]
        point_labels = point_labels[2:]
        onnx_box_coords = box_points
        onnx_box_labels = np.array([2,3])
        input_point = np.array(points, dtype=np.float32)
        input_label = np.array(point_labels, dtype=np.int32)
        onnx_coord = np.concatenate([input_point, onnx_box_coords], axis=0)[None, :, :]
        onnx_label = np.concatenate([input_label, onnx_box_labels], axis=0)[None, :].astype(np.float32)

    scale, new_height, new_width = _compute_scale_to_resize_image(
        image_size=image_size, image=image
    )
    onnx_coord = (
        onnx_coord.astype(float)
        * (new_width / image.shape[1], new_height / image.shape[0])
    ).astype(np.float32)

    onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)
    onnx_has_mask_input = np.array([-1], dtype=np.float32)

    decoder_inputs = {
        "image_embeddings": image_embedding,
        "point_coords": onnx_coord,
        "point_labels": onnx_label,
        "mask_input": onnx_mask_input,
        "has_mask_input": onnx_has_mask_input,
        "orig_im_size": np.array(image.shape[:2], dtype=np.float32),
    }

    masks, _, _ = decoder_session.run(None, decoder_inputs)
    mask = masks[0, 0]  # (1, 1, H, W) -> (H, W)
    mask = mask > 0.0
    if 0:
        imgviz.io.imsave(
            "mask.jpg", imgviz.label2rgb(mask, imgviz.rgb2gray(image))
        )

    contours = skimage.measure.find_contours(np.pad(mask, pad_width=1))
    contour = max(contours, key=_get_contour_length)
    polygon = skimage.measure.approximate_polygon(
        coords=contour,
        tolerance=np.ptp(contour, axis=0).max() / 100,
    )
    polygon = np.clip(polygon, (0, 0), (mask.shape[0] - 1, mask.shape[1] - 1))
    polygon = polygon[:-1]  # drop last point that is duplicate of first point
    if 0:
        image_pil = PIL.Image.fromarray(image)
        imgviz.draw.line_(image_pil, yx=polygon, fill=(0, 255, 0))
        for point in polygon:
            imgviz.draw.circle_(
                image_pil, center=point, diameter=10, fill=(0, 255, 0)
            )
        imgviz.io.imsave("contour.jpg", np.asarray(image_pil))

    return polygon[:, ::-1]  # yx -> xy

Basically, when there are 3 or more points in the points list, I'm just using the first two points as the top-left corner and bottom-right corner of the box. The rest of the points are used as positive/negative point prompts as usual. Doing this gets the job done, but there is no visualization for the box.

Ideally, I think that the points and box should be created using different buttons:
Left Mouse Button: Positive Point
Shift + Left Mouse Button: Negative Point
Middle Mouse Button (On Button Press): Set Top-Left Corner of Box
Middle Mouse Button (On Button Release): Set Bottom-Right Corner of Box

Can someone implement this with a visualization for the box?