Update grounded_sam_simple_demo.py (#308)

Added DEVICE acceleration
IDEA-Research · Jun 17, 2023 · bb8ce18 · bb8ce18
1 parent 0c29c82
commit bb8ce18
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/grounded_sam_simple_demo.py b/grounded_sam_simple_demo.py
@@ -23,6 +23,7 @@
 
 # Building SAM Model and SAM Predictor
 sam = sam_model_registry[SAM_ENCODER_VERSION](checkpoint=SAM_CHECKPOINT_PATH)
+sam.to(device=DEVICE)
 sam_predictor = SamPredictor(sam)
 
 
@@ -103,4 +104,4 @@ def segment(sam_predictor: SamPredictor, image: np.ndarray, xyxy: np.ndarray) ->
 annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)
 
 # save the annotated grounded-sam image
-cv2.imwrite("grounded_sam_annotated_image.jpg", annotated_image)
+cv2.imwrite("grounded_sam_annotated_image.jpg", annotated_image)