jens
updated path to weights
0ef8343
raw
history blame
852 Bytes
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()