import gradio as gr
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
import supervision as sv



def snap(image, video):
    MODEL_TYPE = "vit_b"
    checkpoint = "sam_vit_b_01ec64.pth"
    sam = sam_model_registry[MODEL_TYPE](checkpoint=checkpoint)
    mask_generator = SamAutomaticMaskGenerator(sam)
    #mask_generator = SamAutomaticMaskGenerator(sam, points_per_side=50)
    sam_result = mask_generator.generate(image)
    mask_annotator = sv.MaskAnnotator()
    detections = sv.Detections.from_sam(sam_result=sam_result)
    annotated_image = mask_annotator.annotate(scene=image.copy(), detections=detections)
    return [annotated_image, video]


demo = gr.Interface(
    snap,
    [gr.Image(source="webcam", tool=None), gr.Video(source="webcam")],
    ["image", "video"],
)

if __name__ == "__main__":
    demo.launch()