import argparse
import functools
import json
import os
import sys
import tempfile
import cv2
import numpy as np
import supervision as sv
from groundingdino.util.inference import Model as DinoModel
from imutils import paths
from PIL import Image
from segment_anything import sam_model_registry
from segment_anything import SamAutomaticMaskGenerator
from segment_anything import SamPredictor
from supervision.detection.utils import xywh_to_xyxy
from tqdm import tqdm
sys.path.append("tag2text")
from tag2text.models import tag2text
from config import *
from utils import detect, download_file_hf, segment, generate_tags, show_anns_sv
def process(
tag2text_model,
grounding_dino_model,
sam_predictor,
sam_automask_generator,
image_path,
task,
prompt,
box_threshold,
text_threshold,
iou_threshold,
kernel_size=2,
expand_mask=False,
device="cuda",
output_dir=None,
save_ann=True,
save_mask=False,
):
detections = None
metadata = {"image": {}, "annotations": [], "assets": {}}
if save_mask:
metadata["assets"]["intermediate_mask"] = []
try:
# Load image
image = Image.open(image_path)
image_pil = image.convert("RGB")
image = np.array(image_pil)
orig_image = image.copy()
# Extract image metadata
filename = os.path.basename(image_path)
basename = os.path.splitext(filename)[0]
h, w = image.shape[:2]
metadata["image"]["file_name"] = filename
metadata["image"]["width"] = w
metadata["image"]["height"] = h
# Generate tags
if task in ["auto", "detection"] and prompt == "":
tags, caption = generate_tags(tag2text_model, image_pil, "None", device)
prompt = " . ".join(tags)
# print(f"Caption: {caption}")
# print(f"Tags: {tags}")
# ToDo: Extract metadata
metadata["image"]["caption"] = caption
metadata["image"]["tags"] = tags
if prompt:
metadata["prompt"] = prompt
# Detect boxes
if prompt != "":
detections, phrases, classes = detect(
grounding_dino_model,
image,
caption=prompt,
box_threshold=box_threshold,
text_threshold=text_threshold,
iou_threshold=iou_threshold,
post_process=True,
)
# Save detection image
if output_dir and save_ann:
# Draw boxes
box_annotator = sv.BoxAnnotator()
labels = [
f"{phrases[i]} {detections.confidence[i]:0.2f}"
for i in range(len(phrases))
]
box_image = box_annotator.annotate(
scene=image, detections=detections, labels=labels
)
box_image_path = os.path.join(output_dir, basename + "_detect.png")
metadata["assets"]["detection"] = box_image_path
Image.fromarray(box_image).save(box_image_path)
# Segmentation
if task in ["auto", "segment"]:
kernel = cv2.getStructuringElement(
cv2.MORPH_ELLIPSE, (2 * kernel_size + 1, 2 * kernel_size + 1)
)
if detections:
masks, scores = segment(
sam_predictor, image=orig_image, boxes=detections.xyxy
)
if expand_mask:
masks = [
cv2.dilate(mask.astype(np.uint8), kernel) for mask in masks
]
else:
masks = [
cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, kernel)
for mask in masks
]
detections.mask = masks
binary_mask = functools.reduce(
lambda x, y: x + y, detections.mask
).astype(np.bool)
else:
masks = sam_automask_generator.generate(orig_image)
sorted_generated_masks = sorted(
masks, key=lambda x: x["area"], reverse=True
)
xywh = np.array([mask["bbox"] for mask in sorted_generated_masks])
scores = np.array(
[mask["predicted_iou"] for mask in sorted_generated_masks]
)
if expand_mask:
mask = np.array(
[
cv2.dilate(mask["segmentation"].astype(np.uint8), kernel)
for mask in sorted_generated_masks
]
)
else:
mask = np.array(
[mask["segmentation"] for mask in sorted_generated_masks]
)
detections = sv.Detections(
xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask
)
binary_mask = None
# Save annotated image
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)
mask_image_path = os.path.join(output_dir, basename + "_mask.png")
metadata["assets"]["mask"] = mask_image_path
Image.fromarray(mask_image).save(mask_image_path)
# Save annotation encoding from https://github.com/LUSSeg/ImageNet-S
mask_enc_path = os.path.join(output_dir, basename + "_mask_enc.npy")
np.save(mask_enc_path, res)
metadata["assets"]["mask_enc"] = mask_enc_path
if binary_mask is not None:
cutout_image = np.expand_dims(binary_mask, axis=-1) * orig_image
cutout_image_path = os.path.join(
output_dir, basename + "_cutout.png"
)
Image.fromarray(cutout_image).save(cutout_image_path)
annotated_image_path = os.path.join(
output_dir, basename + "_annotate.png"
)
metadata["assets"]["annotate"] = annotated_image_path
Image.fromarray(annotated_image).save(annotated_image_path)
# ToDo: Extract metadata
if detections:
i = 0
for (xyxy, mask, confidence, _, _), area, box_area in zip(
detections, detections.area, detections.box_area
):
annotation = {
"id": i + 1,
"bbox": [int(x) for x in xyxy],
"box_area": float(box_area),
}
if confidence:
annotation["confidence"] = float(confidence)
annotation["label"] = phrases[i]
if mask is not None:
# annotation["segmentation"] = mask_to_polygons(mask)
annotation["area"] = int(area)
annotation["predicted_iou"] = float(scores[i])
metadata["annotations"].append(annotation)
i += 1
if output_dir and save_mask:
mask_image_path = os.path.join(
output_dir, f"{basename}_mask_{id}.png"
)
metadata["assets"]["intermediate_mask"].append(mask_image_path)
Image.fromarray(mask * 255).save(mask_image_path)
if output_dir:
meta_file_path = os.path.join(output_dir, basename + "_meta.json")
with open(meta_file_path, "w") as fp:
json.dump(metadata, fp)
else:
meta_file = tempfile.NamedTemporaryFile(delete=False, suffix=".json")
meta_file_path = meta_file.name
return meta_file_path
except Exception as error:
raise ValueError(f"global exception: {error}")
def main(args: argparse.Namespace) -> None:
device = args.device
prompt = args.prompt
task = args.task
tag2text_model = None
grounding_dino_model = None
sam_predictor = None
sam_automask_generator = 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_type
tag2text_checkpoint = os.path.join(
abs_weight_dir, tag2text_dict[tag2text_type]["checkpoint_file"]
)
if not os.path.exists(tag2text_checkpoint):
print(f"Downloading weights for Tag2Text {tag2text_type} model")
os.system(
f"wget {tag2text_dict[tag2text_type]['checkpoint_url']} -O {tag2text_checkpoint}"
)
tag2text_model = tag2text.tag2text_caption(
pretrained=tag2text_checkpoint,
image_size=384,
vit="swin_b",
delete_tag_index=delete_tag_index,
)
# threshold for tagging
# we reduce the threshold to obtain more tags
tag2text_model.threshold = 0.64
tag2text_model.to(device)
tag2text_model.eval()
if task in ["auto", "detection"] or prompt != "":
print("Loading Grounding Dino model...")
dino_type = args.dino_type
dino_checkpoint = os.path.join(
abs_weight_dir, dino_dict[dino_type]["checkpoint_file"]
)
dino_config_file = os.path.join(
abs_weight_dir, dino_dict[dino_type]["config_file"]
)
if not os.path.exists(dino_checkpoint):
print(f"Downloading weights for Grounding Dino {dino_type} model")
dino_repo_id = dino_dict[dino_type]["repo_id"]
download_file_hf(
repo_id=dino_repo_id,
filename=dino_dict[dino_type]["checkpoint_file"],
cache_dir=weight_dir,
)
download_file_hf(
repo_id=dino_repo_id,
filename=dino_dict[dino_type]["checkpoint_file"],
cache_dir=weight_dir,
)
grounding_dino_model = DinoModel(
model_config_path=dino_config_file,
model_checkpoint_path=dino_checkpoint,
device=device,
)
if task in ["auto", "segment"]:
print("Loading SAM...")
sam_type = args.sam_type
sam_checkpoint = os.path.join(
abs_weight_dir, sam_dict[sam_type]["checkpoint_file"]
)
if not os.path.exists(sam_checkpoint):
print(f"Downloading weights for SAM {sam_type}")
os.system(
f"wget {sam_dict[sam_type]['checkpoint_url']} -O {sam_checkpoint}"
)
sam = sam_model_registry[sam_type](checkpoint=sam_checkpoint)
sam.to(device=device)
sam_predictor = SamPredictor(sam)
sam_automask_generator = SamAutomaticMaskGenerator(sam)
if not os.path.exists(args.input):
raise ValueError("The input directory doesn't exist!")
elif not os.path.isdir(args.input):
image_paths = [args.input]
else:
image_paths = paths.list_images(args.input)
os.makedirs(args.output, exist_ok=True)
with tqdm(image_paths) as pbar:
for image_path in pbar:
pbar.set_postfix_str(f"Processing {image_path}")
process(
tag2text_model=tag2text_model,
grounding_dino_model=grounding_dino_model,
sam_predictor=sam_predictor,
sam_automask_generator=sam_automask_generator,
image_path=image_path,
task=task,
prompt=prompt,
box_threshold=box_threshold,
text_threshold=text_threshold,
iou_threshold=iou_threshold,
device=device,
output_dir=args.output,
save_ann=save_ann,
save_mask=save_mask,
)
if __name__ == "__main__":
if not os.path.exists(abs_weight_dir):
os.makedirs(abs_weight_dir, exist_ok=True)
parser = argparse.ArgumentParser(
description=(
"Runs automatic detection and mask generation on an input image or directory of images"
)
)
parser.add_argument(
"--input",
"-i",
type=str,
required=True,
help="Path to either a single input image or folder of images.",
)
parser.add_argument(
"--output",
"-o",
type=str,
required=True,
help="Path to the directory where masks will be output.",
)
parser.add_argument(
"--sam-type",
type=str,
default=default_sam,
choices=sam_dict.keys(),
help="The type of SA model use for segmentation.",
)
parser.add_argument(
"--tag2text-type",
type=str,
default=default_tag2text,
choices=tag2text_dict.keys(),
help="The type of Tag2Text model use for tags and caption generation.",
)
parser.add_argument(
"--dino-type",
type=str,
default=default_dino,
choices=dino_dict.keys(),
help="The type of Grounding Dino model use for promptable object detection.",
)
parser.add_argument(
"--task",
help="Task to run",
default="auto",
choices=["auto", "detect", "segment"],
type=str,
)
parser.add_argument(
"--prompt",
help="Detection prompt",
default="",
type=str,
)
parser.add_argument(
"--box-threshold", type=float, default=0.25, help="box threshold"
)
parser.add_argument(
"--text-threshold", type=float, default=0.2, help="text threshold"
)
parser.add_argument(
"--iou-threshold", type=float, default=0.5, help="iou threshold"
)
parser.add_argument(
"--kernel-size",
type=int,
default=2,
choices=range(1, 6),
help="kernel size use for smoothing/expanding segment masks",
)
parser.add_argument(
"--expand-mask",
action="store_true",
default=False,
help="If True, expanding segment masks for smoother output.",
)
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",
default=False,
help="If True, save all intermidiate masks.",
)
parser.add_argument(
"--device", type=str, default="cuda", help="The device to run generation on."
)
args = parser.parse_args()
main(args)