File size: 852 Bytes
5c0b534
9b4ee8f
 
 
5c0b534
 
 
9b4ee8f
0ef8343
9b4ee8f
 
 
 
 
 
 
 
5c0b534
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
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()