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
import PIL.Image
import cv2

import gradio as gr


TITLE = """
# YOLO-World-Seg

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

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", "horse",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()

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()
    with open("input.jpeg", "wb") as f:
        f.write(input_image)
        
    class_names = [class_name.strip() for class_name in class_names.split(',')]
    
    texts = [[t.strip()] for t in class_names.split(",")] + [[" "]]
    data_info = runner.pipeline(dict(img_id=0, img_path="input.jpeg",
                                     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 = box_annotator.annotate(input_image, 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."
    ))

with gr.Blocks() as demo:
    gr.Markdown(TITLE)
    with gr.Accordion("Configuration", open=False):
        confidence_threshold_component.render()
        iou_threshold_component.render()
    with gr.Tab(label="Image"):
        with gr.Row():
            input_image_component = gr.Image(
                type='numpy',
                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,
            ],
            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,
        ],
        outputs=output_image_component
    )
    


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