update script

annotate_anything.py

@@ -33,8 +33,9 @@ def process(
     box_threshold,
     text_threshold,
     iou_threshold,
-    device,
+    device="cuda",
     output_dir=None,
+    save_ann=True,
     save_mask=False,
 ):
     detections = None
@@ -84,7 +85,7 @@ def process(
             )
 
             # Save detection image
-            if output_dir:
+            if output_dir and save_ann:
                 # Draw boxes
                 box_annotator = sv.BoxAnnotator()
                 labels = [
@@ -123,7 +124,7 @@ def process(
                 )
 
             # Save annotated image
-            if output_dir:
+            if output_dir and save_ann:
                 mask_annotator = sv.MaskAnnotator()
                 mask_image, res = show_anns_sv(detections)
                 annotated_image = mask_annotator.annotate(image, detections=detections)
@@ -197,12 +198,13 @@ def main(args: argparse.Namespace) -> None:
     box_threshold = args.box_threshold
     text_threshold = args.text_threshold
     iou_threshold = args.iou_threshold
+    save_ann = not args.no_save_ann
     save_mask = args.save_mask
 
     # load model
     if task in ["auto", "detection"] and prompt == "":
         print("Loading Tag2Text model...")
-        tag2text_type = args.tag2text
+        tag2text_type = args.tag2text_type
         tag2text_checkpoint = os.path.join(
             abs_weight_dir, tag2text_dict[tag2text_type]["checkpoint_file"]
         )
@@ -225,7 +227,7 @@ def main(args: argparse.Namespace) -> None:
 
     if task in ["auto", "detection"] or prompt != "":
         print("Loading Grounding Dino model...")
-        dino_type = args.dino
+        dino_type = args.dino_type
         dino_checkpoint = os.path.join(
             abs_weight_dir, dino_dict[dino_type]["checkpoint_file"]
         )
@@ -253,7 +255,7 @@ def main(args: argparse.Namespace) -> None:
 
     if task in ["auto", "segment"]:
         print("Loading SAM...")
-        sam_type = args.sam
+        sam_type = args.sam_type
         sam_checkpoint = os.path.join(
             abs_weight_dir, sam_dict[sam_type]["checkpoint_file"]
         )
@@ -292,6 +294,7 @@ def main(args: argparse.Namespace) -> None:
                 iou_threshold=iou_threshold,
                 device=device,
                 output_dir=args.output,
+                save_ann=save_ann,
                 save_mask=save_mask,
             )
 
@@ -319,34 +322,31 @@ if __name__ == "__main__":
         "-o",
         type=str,
         required=True,
-        help=(
-            "Path to the directory where masks will be output. Output will be either a folder "
-            "of PNGs per image or a single json with COCO-style masks."
-        ),
+        help="Path to the directory where masks will be output.",
     )
 
     parser.add_argument(
-        "--sam",
+        "--sam-type",
         type=str,
         default=default_sam,
         choices=sam_dict.keys(),
-        help="The type of SA model to load",
+        help="The type of SA model use for segmentation.",
     )
 
     parser.add_argument(
-        "--tag2text",
+        "--tag2text-type",
         type=str,
         default=default_tag2text,
         choices=tag2text_dict.keys(),
-        help="The path to the Tag2Text checkpoint to use for tags and caption generation.",
+        help="The type of Tag2Text model use for tags and caption generation.",
     )
 
     parser.add_argument(
-        "--dino",
+        "--dino-type",
         type=str,
         default=default_dino,
         choices=dino_dict.keys(),
-        help="The config file of Grounding Dino model to load",
+        help="The type of Grounding Dino model use for promptable object detection.",
     )
 
     parser.add_argument(
@@ -373,6 +373,12 @@ if __name__ == "__main__":
         "--iou-threshold", type=float, default=0.5, help="iou threshold"
     )
 
+    parser.add_argument(
+        "--no-save-ann",
+        action="store_true",
+        default=False,
+        help="If False, save original image with blended masks and detection boxes.",
+    )
     parser.add_argument(
         "--save-mask",
         action="store_true",