app.py

import os
os.system("mim install 'mmengine>=0.7.0'")
os.system("mim install mmcv")
os.system("mim install 'mmdet>=3.0.0'") 
os.system("pip install -e .")


import numpy as np
import torch
from mmengine.config import Config
from mmengine.dataset import Compose
from mmengine.runner import Runner
from mmengine.runner.amp import autocast
from mmyolo.registry import RUNNERS
from torchvision.ops import nms
import supervision as sv
from PIL import Image
import cv2
import spaces

import gradio as gr


TITLE = """
# YOLO-World-Seg

This is a demo of zero-shot object detection and instance segmentation using only
[YOLO-World](https://github.com/AILab-CVC/YOLO-World) done via newest model YOLO-World-Seg.

Annototions Powered by [Supervision](https://github.com/roboflow/supervision).
"""

EXAMPLES = [
    ["https://media.roboflow.com/efficient-sam/corgi.jpg", "dog",0.5,0.5,0.5,100],
    ["https://media.roboflow.com/efficient-sam/horses.jpg", "horses",0.5,0.5,0.5,100],
    ["https://media.roboflow.com/efficient-sam/bears.jpg", "bear",0.5,0.5,0.5,100],
]

box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
mask_annotator = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)

def load_runner():
    cfg = Config.fromfile(
    "./configs/segmentation/yolo_world_seg_l_dual_vlpan_2e-4_80e_8gpus_seghead_finetune_lvis.py"
    )
    cfg.work_dir = "."
    cfg.load_from = "yolo_world_seg_l_dual_vlpan_2e-4_80e_8gpus_seghead_finetune_lvis-5a642d30.pth"
    runner = Runner.from_cfg(cfg)
    runner.call_hook("before_run")
    runner.load_or_resume()
    pipeline = cfg.test_dataloader.dataset.pipeline
    runner.pipeline = Compose(pipeline)
    runner.model.eval()
    return runner

@spaces.GPU
def run_image(
        input_image,
        class_names="person,car,bus,truck",
        score_thr=0.05,
        iou_thr=0.5,
        nms_thr=0.5,
        max_num_boxes=100,
        ):
    runner = load_runner()
    
    image_path='./work_dirs/input.png'
    os.makedirs('./work_dirs', exist_ok=True)
    input_image.save(image_path)
            
    texts = [[t.strip()] for t in class_names.split(",")] + [[" "]]
    data_info = runner.pipeline(dict(img_id=0, img_path=image_path,
                                     texts=texts))

    data_batch = dict(
        inputs=data_info["inputs"].unsqueeze(0),
        data_samples=[data_info["data_samples"]],
    )

    with autocast(enabled=False), torch.no_grad():
        output = runner.model.test_step(data_batch)[0]
        runner.model.class_names = texts
        pred_instances = output.pred_instances

    keep_idxs = nms(pred_instances.bboxes, pred_instances.scores, iou_threshold=iou_thr)
    pred_instances = pred_instances[keep_idxs]
    pred_instances = pred_instances[pred_instances.scores.float() > score_thr]

    if len(pred_instances.scores) > max_num_boxes:
        indices = pred_instances.scores.float().topk(max_num_boxes)[1]
        pred_instances = pred_instances[indices]
    output.pred_instances = pred_instances
    result = pred_instances.cpu().numpy()
    detections = sv.Detections(
    xyxy=result['bboxes'],
    class_id=result['labels'],
    confidence=result['scores'],
    mask = result['masks']
    )
    detections = detections.with_nms(threshold=nms_thr)

    labels = [
        f"{class_id} {confidence:.3f}"
        for class_id, confidence
        in zip(detections.class_id, detections.confidence)
    ]
    
    svimage = np.array(input_image)
    svimage = box_annotator.annotate(svimage, detections)
    svimage = label_annotator.annotate(svimage, detections, labels)
    svimage = mask_annotator.annotate(svimage,detections)
    return svimage

confidence_threshold_component = gr.Slider(
    minimum=0,
    maximum=1.0,
    value=0.3,
    step=0.01,
    label="Confidence Threshold",
    info=(
        "The confidence threshold for the YOLO-World model. Lower the threshold to "
        "reduce false negatives, enhancing the model's sensitivity to detect "
        "sought-after objects. Conversely, increase the threshold to minimize false "
        "positives, preventing the model from identifying objects it shouldn't."
    ))

iou_threshold_component = gr.Slider(
    minimum=0,
    maximum=1.0,
    value=0.5,
    step=0.01,
    label="IoU Threshold",
    info=(
        "The Intersection over Union (IoU) threshold for non-maximum suppression. "
        "Decrease the value to lessen the occurrence of overlapping bounding boxes, "
        "making the detection process stricter. On the other hand, increase the value "
        "to allow more overlapping bounding boxes, accommodating a broader range of "
        "detections."
    ))

nms_threshold_component = gr.Slider(
    minimum=0,
    maximum=1.0,
    value=0.5,
    step=0.01,
    label="NMS Threshold",
    info=(
        "The Non-Maximum Suppression (NMS) Threshold is a parameter that determines the Intersection over Union (IoU) threshold for suppressing bounding boxes. "
        "A lower value will reduce the likelihood of overlapping bounding boxes, resulting in a more stringent detection process. Conversely, a higher value "
        "will permit more overlapping bounding boxes, thereby allowing for a wider variety of detections."
    ))

with gr.Blocks() as demo:
    gr.Markdown(TITLE)
    with gr.Accordion("Configuration", open=False):
        confidence_threshold_component.render()
        iou_threshold_component.render()
        nms_threshold_component.render()
    with gr.Tab(label="Image"):
        with gr.Row():
            input_image_component = gr.Image(
                type='pil',
                label='Input Image'
            )
            output_image_component = gr.Image(
                type='numpy',
                label='Output Image'
            )
        with gr.Row():
            image_categories_text_component = gr.Textbox(
                label='Categories',
                placeholder='comma separated list of categories',
                scale=7
            )
            image_submit_button_component = gr.Button(
                value='Submit',
                scale=1,
                variant='primary'
            )
        gr.Examples(
            fn=run_image,
            examples=EXAMPLES,
            inputs=[
                input_image_component,
                image_categories_text_component,
                confidence_threshold_component,
                iou_threshold_component,
                nms_threshold_component
            ],
            outputs=output_image_component
        )


    image_submit_button_component.click(
        fn=run_image,
        inputs=[
            input_image_component,
            image_categories_text_component,
            confidence_threshold_component,
            iou_threshold_component,
            nms_threshold_component
        ],
        outputs=output_image_component
    )
    


demo.launch(debug=False, show_error=True)