🚀 [New] Huggingface examples!

README.md

@@ -1,6 +1,8 @@
 # YOLO: Official Implementation of YOLOv9, YOLOv7
 
+![GitHub License](https://img.shields.io/github/license/WongKinYiu/YOLO)
 ![WIP](https://img.shields.io/badge/status-WIP-orange)
+[![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/henry000/YOLO)
 > [!IMPORTANT]
 > This project is currently a Work In Progress and may undergo significant changes. It is not recommended for use in production environments until further notice. Please check back regularly for updates.
 >
@@ -35,7 +37,7 @@ pip install -r requirements.txt
 
 | Tools | pip 🐍 | HuggingFace 🤗 | Docker 🐳 |
 | -------------------- | :----: | :--------------: | :-------: |
-| Compatibility       | ✅     | ❔               | 🧪        |
+| Compatibility       | ✅     | ✅               | 🧪        |
 
 |  Phase    | Training | Validation | Inference |
 | ------------------- | :------: | :---------: | :-------: |

demo/hf_demo.py

@@ -0,0 +1,76 @@
+import sys
+from pathlib import Path
+
+import gradio
+import torch
+from omegaconf import OmegaConf
+
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+
+from yolo import (
+    AugmentationComposer,
+    NMSConfig,
+    Vec2Box,
+    bbox_nms,
+    create_model,
+    draw_bboxes,
+)
+
+DEFAULT_MODEL = "v9-c"
+IMAGE_SIZE = (640, 640)
+
+
+def load_model(model_name, device):
+    model_cfg = OmegaConf.load(f"yolo/config/model/{model_name}.yaml")
+    model_cfg.model.auxiliary = {}
+    model = create_model(model_cfg, True)
+    model.to(device).eval()
+    return model
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = load_model(DEFAULT_MODEL, device)
+v2b = Vec2Box(model, IMAGE_SIZE, device)
+class_list = OmegaConf.load("yolo/config/general.yaml").class_list
+
+transform = AugmentationComposer([])
+
+
+def predict(model_name, image, nms_confidence, nms_iou):
+    global DEFAULT_MODEL, model, device, v2b, class_list
+    if model_name != DEFAULT_MODEL:
+        model = load_model(model_name, device)
+        v2b = Vec2Box(model, IMAGE_SIZE, device)
+        DEFAULT_MODEL = model_name
+
+    image_tensor, _, rev_tensor = transform(image)
+
+    image_tensor = image_tensor.to(device)[None]
+    rev_tensor = rev_tensor.to(device)
+
+    with torch.no_grad():
+        predict = model(image_tensor)
+        pred_class, _, pred_bbox = v2b(predict["Main"])
+
+    nms_config = NMSConfig(nms_confidence, nms_iou)
+
+    pred_bbox = pred_bbox / rev_tensor[0] - rev_tensor[None, None, 1:]
+    pred_bbox = bbox_nms(pred_class, pred_bbox, nms_config)
+    result_image = draw_bboxes(image, pred_bbox, idx2label=class_list)
+
+    return result_image
+
+
+interface = gradio.Interface(
+    fn=predict,
+    inputs=[
+        gradio.components.Dropdown(choices=["v9-c", "v9-m", "v9-s"], value="v9-c", label="Model Name"),
+        gradio.components.Image(type="pil", label="Input Image"),
+        gradio.components.Slider(0, 1, step=0.01, value=0.5, label="NMS Confidence Threshold"),
+        gradio.components.Slider(0, 1, step=0.01, value=0.5, label="NMS IoU Threshold"),
+    ],
+    outputs=gradio.components.Image(type="pil", label="Output Image"),
+)
+
+if __name__ == "__main__":
+    interface.launch()

yolo/__init__.py

@@ -1,4 +1,4 @@
-from yolo.config.config import Config
+from yolo.config.config import Config, NMSConfig
 from yolo.model.yolo import create_model
 from yolo.tools.data_loader import AugmentationComposer, create_dataloader
 from yolo.tools.drawer import draw_bboxes
@@ -10,6 +10,7 @@ from yolo.utils.logging_utils import custom_logger
 all = [
     "create_model",
     "Config",
+    "NMSConfig",
     "custom_logger",
     "validate_log_directory",
     "draw_bboxes",