annotate_anything.py

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)