Add files

ExportModel.py

@@ -0,0 +1,40 @@
+import torch
+import os
+from copy import deepcopy
+
+class ModelExporter(torch.nn.Module):
+    def __init__(self, yoloModel, device='cpu'):
+        super(ModelExporter, self).__init__()
+        model = deepcopy(yoloModel).to(device)
+        for p in model.parameters():
+            p.requires_grad = False
+        model.eval()
+        model.float()
+        model = model.fuse()
+
+        self.model = model
+        self.device = device
+
+    def forward(self, x, txt_feats):
+        return self.model.predict(x, txt_feats=txt_feats)
+
+    def export(self, output_dir, model_name, img_width, img_height, num_classes):
+        x = torch.randn(1, 3, img_width, img_height, requires_grad=False).to(self.device)
+        txt_feats = torch.randn(1, num_classes, 512, requires_grad=False).to(self.device)
+
+        print(x.shape, txt_feats.shape)
+
+        # Export model
+        onnx_name = model_name + ".onnx"
+        os.makedirs(output_dir, exist_ok=True)
+        output_path = f"{output_dir}/{onnx_name}"
+        with torch.no_grad():
+            torch.onnx.export(self,
+                              (x, txt_feats),
+                              output_path,
+                              do_constant_folding=True,
+                              opset_version=17,
+                              input_names=["images", "txt_feats"],
+                              output_names=["output"])
+
+        return output_path

app.py

@@ -0,0 +1,40 @@
+import gradio as gr
+from ultralytics import YOLOWorld
+
+from ExportModel import ModelExporter
+
+model_list = ['yolov8s-worldv2', 'yolov8m-worldv2', 'yolov8l-worldv2', 'yolov8x-worldv2']
+
+def export_model(model, width, height, num_classes):
+    model_name = model
+    img_width = width
+    img_height = height
+    num_classes = num_classes
+
+    yoloModel = YOLOWorld(model_name)
+    yoloModel.set_classes([""] * num_classes)
+
+    # Initialize model exporter
+    export_model = ModelExporter(yoloModel.model)
+
+    # Export model
+    output_path = export_model.export("temp", model_name, img_width, img_height, num_classes)
+
+    return output_path
+
+
+demo = gr.Interface(
+    export_model,
+    [
+        gr.Dropdown(model_list, label="model", value=model_list[0]),
+        gr.Slider(32, 4096, step=32, value=640, label="width"),
+        gr.Slider(32, 4096, step=32, value=480, label="height"),
+        gr.Number(label="num_classes", value=1),
+    ],
+    "file",
+    title="ONNX Export Ultralytics YOLO-World Open Vocabulary Object Detection",
+    description="Demo to export Ultralytics YOLO-World Open Vocabulary Object Detection model to ONNX",
+    api_name="export"
+)
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+ultralytics
+gradio