quoccuongle/Grounded-SAM-2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: grounded_sam2_hf_mode cli arguments (#52)

Browse Source
This commit is contained in:
kwikwag
2024-09-24 15:33:11 +03:00
committed by GitHub
parent 81ac531aa9
commit e899ad99e8
1 changed files with 20 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
grounded_sam2_hf_model_demo.py
View File

@@ -1,3 +1,4 @@
import argparse
import os
import cv2
import json
@@ -16,14 +17,25 @@ from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
"""
Hyper parameters
"""
GROUNDING_MODEL = "IDEA-Research/grounding-dino-tiny"
TEXT_PROMPT = "car. tire."
IMG_PATH = "notebooks/images/truck.jpg"
SAM2_CHECKPOINT = "./checkpoints/sam2_hiera_large.pt"
SAM2_MODEL_CONFIG = "sam2_hiera_l.yaml"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
OUTPUT_DIR = Path("outputs/grounded_sam2_hf_model_demo")
DUMP_JSON_RESULTS = True
parser = argparse.ArgumentParser()
parser.add_argument('--grounding-model', default="IDEA-Research/grounding-dino-tiny")
parser.add_argument("--text-prompt", default="car. tire.")
parser.add_argument("--img-path", default="notebooks/images/truck.jpg")
parser.add_argument("--sam2-checkpoint", default="./checkpoints/sam2_hiera_large.pt")
parser.add_argument("--sam2-model-config", default="sam2_hiera_l.yaml")
parser.add_argument("--output-dir", default="outputs/grounded_sam2_hf_model_demo")
parser.add_argument("--no-dump-json", action="store_true")
parser.add_argument("--force-cpu", action="store_true")
args = parser.parse_args()
GROUNDING_MODEL = args.grounding_model
TEXT_PROMPT = args.text_prompt
IMG_PATH = args.img_path
SAM2_CHECKPOINT = args.sam2_checkpoint
SAM2_MODEL_CONFIG = args.sam2_model_config
DEVICE = "cuda" if torch.cuda.is_available() and not args.force_cpu else "cpu"
OUTPUT_DIR = Path(args.output_dir)
DUMP_JSON_RESULTS = not args.no_dump_json
# create output directory
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)