jens commited on
Commit
9b4ee8f
·
1 Parent(s): 5c0b534

implemented SAM segment anything

Browse files
Files changed (2) hide show
  1. app.py +13 -1
  2. requirements.txt +6 -0
app.py CHANGED
@@ -1,8 +1,20 @@
1
  import gradio as gr
 
 
 
2
 
3
 
4
  def snap(image, video):
5
- return [image, video]
 
 
 
 
 
 
 
 
 
6
 
7
 
8
  demo = gr.Interface(
 
1
  import gradio as gr
2
+ from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
3
+ import supervision as sv
4
+
5
 
6
 
7
  def snap(image, video):
8
+ MODEL_TYPE = "vit_b"
9
+ checkpoint = "checkpoints/sam_vit_b_01ec64.pth"
10
+ sam = sam_model_registry[MODEL_TYPE](checkpoint=checkpoint)
11
+ mask_generator = SamAutomaticMaskGenerator(sam)
12
+ #mask_generator = SamAutomaticMaskGenerator(sam, points_per_side=50)
13
+ sam_result = mask_generator.generate(image)
14
+ mask_annotator = sv.MaskAnnotator()
15
+ detections = sv.Detections.from_sam(sam_result=sam_result)
16
+ annotated_image = mask_annotator.annotate(scene=image.copy(), detections=detections)
17
+ return [annotated_image, video]
18
 
19
 
20
  demo = gr.Interface(
requirements.txt CHANGED
@@ -1 +1,7 @@
1
  gradio
 
 
 
 
 
 
 
1
  gradio
2
+ huggingface_hub
3
+ segment-anything
4
+ supervision
5
+ torch
6
+ torchvision
7
+ opencv-python