Port to SAM2.1

[aakash-vwo]

Saw you are assoicated with Roboflow, so I am creating this issue
Could you nudge me in the right direction to port SoM to SAM2.1?

I cloned
https://huggingface.co/spaces/Roboflow/SoM/blob/main/app.py

and rewrote sam_utils.py

import numpy as np
import supervision as sv
import torch
from PIL import Image
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

SAM_CHECKPOINT = "checkpoints/sam2_hiera_small.pt"
SAM_CONFIG = "sam2_hiera_s.yaml"


def load_sam_model(device: torch.device) -> SAM2ImagePredictor:
    model = build_sam2(SAM_CONFIG, SAM_CHECKPOINT, device=device)
    return SAM2ImagePredictor(sam_model=model)


def sam_inference(image: np.ndarray, model: SAM2ImagePredictor) -> sv.Detections:
    pil_image = Image.fromarray(image)
    model.set_image(np.array(pil_image.convert("RGB")))
    masks, _, _ = model.predict()

    return sv.Detections(xyxy=sv.mask_to_xyxy(masks), mask=masks.astype(bool))


def sam_interactive_inference(
    image: np.ndarray, mask: np.ndarray, model: SAM2ImagePredictor
) -> sv.Detections:
    pil_image = Image.fromarray(image)
    model.set_image(np.array(pil_image.convert("RGB")))
    masks = []
    for polygon in sv.mask_to_polygons(mask.astype(bool)):
        random_point_indexes = np.random.choice(polygon.shape[0], size=5, replace=True)
        input_point = polygon[random_point_indexes]
        input_label = np.ones(5)
        mask, _, _ = model.predict(
            point_coords=input_point,
            point_labels=input_label,
            multimask_output=False,
        )
        masks.append(mask[0])  # SAM2 returns a list of masks, we take the first one
    masks = np.array(masks, dtype=bool)
    return sv.Detections(xyxy=sv.mask_to_xyxy(masks), mask=masks)

I think I am doing something wrong because I only get 1 annotation