add onnx export and inference

.gitignore

@@ -1,3 +1,5 @@
 .venv
 __pycache__
-data/cache
+data/cache
+data/onnx
+.DS_Store

README.md

@@ -6,4 +6,4 @@ app_port: 7000
 pinned: true
 ---
 
-# Simple DETR gradio implementation (object detections & panoptic segmentation)
+# Simple DETR gradio implementation (object detection & panoptic segmentation)

detr/{detr.py → detr_models.py}

@@ -10,7 +10,10 @@ from torch import nn
 from torchvision.models import resnet50
 from panopticapi.utils import id2rgb, rgb2id
 from supervision import Detections, BoxAnnotator, MaskAnnotator
+import onnx
+import onnxruntime
 from PIL import Image
+from pathlib import Path
 
 torch.set_grad_enabled(False)
 
@@ -18,11 +21,14 @@ torch.set_grad_enabled(False)
 # https://colab.research.google.com/github/facebookresearch/detr/blob/colab/notebooks/detr_demo.ipynb#scrollTo=cfCcEYjg7y46
 
 DETR_DEMO_WEIGHTS_URI = "https://dl.fbaipublicfiles.com/detr/detr_demo-da2a99e9.pth"
-
 TORCH_HOME = os.path.abspath(os.curdir) + "/data/cache"
-
+ONNX_DIR = os.path.abspath(os.curdir) + "/data/onnx"
 os.environ["TORCH_HOME"] = TORCH_HOME
 
+Path(TORCH_HOME).mkdir(exist_ok=True)
+Path(ONNX_DIR).mkdir(exist_ok=True)
+
+
 print("Torch home:", TORCH_HOME)
 
 
@@ -40,6 +46,17 @@ def normalize_img(image):
     return transform(image).unsqueeze(0)
 
 
+def normalize_img_800_800(image):
+    transform = T.Compose(
+        [
+            T.Resize((800, 800)),
+            T.ToTensor(),
+            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+    return transform(image).unsqueeze(0)
+
+
 # for output bounding box post-processing
 def box_cxcywh_to_xyxy(x):
     x_c, y_c, w, h = x.unbind(1)
@@ -199,6 +216,100 @@ class SimpleDetr:
         )
         return annotated
 
+    def export(self):
+        model_path = f"{ONNX_DIR}/detr_simple_demo_onnx.onnx"
+        dummy_image = torch.ones(1, 3, 800, 800, device="cpu")
+        input_names = ["inputs"]
+        output_names = ["pred_logits", "pred_boxes"]
+        torch.onnx.export(
+            self.model,
+            dummy_image,
+            model_path,
+            input_names=input_names,
+            output_names=output_names,
+            # dynamic_axes={input_names[0]: {0: "batch_size", 2: "height", 3: "width"}}, #!TODO
+            export_params=True,
+            training=torch.onnx.TrainingMode.EVAL,
+            opset_version=14,
+        )
+        onnx_model = onnx.load(model_path)
+
+        # Check the model
+        try:
+            onnx.checker.check_model(onnx_model)
+        except onnx.checker.ValidationError as e:
+            print(f"The model is invalid: {e}")
+        else:
+            print("The model is valid!")
+        return model_path
+
+
+class SimpleDetrOnnx:
+    @cache
+    def __init__(self):
+        self.box_annotator: BoxAnnotator = BoxAnnotator()
+        onnx_sess_opts = onnxruntime.SessionOptions()
+        onnx_sess_opts.graph_optimization_level = (
+            onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+            # onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
+        )
+        onnx_sess_opts.enable_mem_pattern = True
+        onnx_sess_opts.enable_cpu_mem_arena = True
+        self.ort_session = onnxruntime.InferenceSession(
+            f"{ONNX_DIR}/detr_simple_demo.onnx",
+            sess_options=onnx_sess_opts,
+            providers=[
+                "CUDAExecutionProvider",
+                "CoreMLExecutionProvider",
+                "CPUExecutionProvider",
+            ],
+        )
+        self.classes = {}
+        self.metadata = self.ort_session.get_modelmeta()
+        self.providers = self.ort_session.get_providers()
+        print(f"[OnnxRuntime] Providers:{self.providers}")
+        print(
+            f"[OnnxRuntime] medatadata: {self.metadata.custom_metadata_map} {type(self.metadata.custom_metadata_map)}"
+        )
+
+    def detect(self, image, conf):
+        # dummy_image = np.ones((1, 3, 600, 800), dtype=np.float32)
+        im = normalize_img_800_800(image).numpy()
+        print("SHAPE", im.shape)
+        ort_inputs = {self.ort_session.get_inputs()[0].name: im}
+        outputs = self.ort_session.run(None, ort_inputs)
+        pred_logits = torch.tensor(
+            outputs[0]
+        )  # conversion to torch for simplicity (softmax etc)
+        pred_boxes = torch.tensor(outputs[1])
+        scores = pred_logits.softmax(-1)[0, :, :-1]
+        keep = scores.max(-1).values > conf
+        bboxes_scaled = rescale_bboxes(pred_boxes[0, keep], image.size)
+        probas = scores[keep]
+        class_id = []
+        confidence = []
+        for prob in probas:
+            cls_id = prob.argmax()
+            c = prob[cls_id]
+            class_id.append(int(cls_id))
+            confidence.append(float(c))
+        print(class_id, confidence)
+        detections = Detections(
+            xyxy=bboxes_scaled.cpu().detach().numpy(),
+            class_id=np.array(class_id),
+            confidence=np.array(confidence),
+        )
+        annotated = self.box_annotator.annotate(
+            scene=np.array(image),
+            skip_label=False,
+            detections=detections,
+            labels=[
+                f"{CLASSES[cls_id]} {conf:.2f}"
+                for cls_id, conf in zip(detections.class_id, detections.confidence)
+            ],
+        )
+        return annotated
+
 
 class PanopticDetrResenet101:
     @cache
@@ -235,6 +346,32 @@ class PanopticDetrResenet101:
             panoptic_seg[panoptic_seg_id == id] = np.asarray(next(palette)) * 255
         return panoptic_seg
 
+    def export(self):
+        model_path = f"{ONNX_DIR}/detr_resnet101_panoptic.onnx"
+        dummy_image = torch.ones(1, 3, 800, 800, device="cpu")
+        input_names = ["inputs"]
+        output_names = ["pred_logits", "pred_boxes", "pred_masks"]
+        torch.onnx.export(
+            self.model,
+            dummy_image,
+            model_path,
+            input_names=input_names,
+            output_names=output_names,
+            export_params=True,
+            training=torch.onnx.TrainingMode.EVAL,
+            opset_version=14,
+        )
+        onnx_model = onnx.load(model_path)
+
+        # Check the model
+        try:
+            onnx.checker.check_model(onnx_model)
+        except onnx.checker.ValidationError as e:
+            print(f"The model is invalid: {e}")
+        else:
+            print("The model is valid!")
+        return model_path
+
 
 # COCO classes
 CLASSES = [
@@ -330,3 +467,6 @@ CLASSES = [
     "hair drier",
     "toothbrush",
 ]
+
+# model = SimpleDetr()
+# model.export()

detr/main_gradio.py

@@ -3,20 +3,26 @@ import supervision as sv
 import os
 from time import perf_counter
 
-from detr import SimpleDetr, PanopticDetrResenet101
+from detr_models import SimpleDetr, PanopticDetrResenet101, SimpleDetrOnnx, ONNX_DIR
 
+IMAGES_DIR = os.path.abspath(os.curdir) + "/data/images"
 ASSETS_DIR = os.path.abspath(os.curdir) + "/data/assets"
-
-print("Assets:", ASSETS_DIR)
+print("images:", IMAGES_DIR)
 
 
 def run_inference(image, confidence, model_name, progress=gr.Progress(track_tqdm=True)):
     progress(0.1, "loading model..")
+    if not image:
+        raise gr.Error("Provide image.")
     t0 = perf_counter()
-    if model_name == "detr_demo_boxes":
+    if model_name == "detr_simple_demo":
         model = SimpleDetr()
-    else:
+    elif model_name == "detr_resnet101_panoptic":
         model = PanopticDetrResenet101()
+    elif model_name == "detr_simple_demo_onnx":
+        if not os.path.exists(f"{ONNX_DIR}/detr_simple_demo_onnx.onnx"):
+            raise gr.Error("ONNX model not found, please export it first!")
+        model = SimpleDetrOnnx()
     t1 = perf_counter()
     progress(0.1, "Inference..")
 
@@ -25,43 +31,111 @@ def run_inference(image, confidence, model_name, progress=gr.Progress(track_tqdm
     return annotated_img, {"load_model": t1 - t0, "inference": t2 - t1}, None
 
 
-with gr.Blocks() as inference_gradio:
-    gr.Markdown("# DETR inference")
-    with gr.Row():
-        with gr.Column():
-            img_file = gr.Image(type="pil")
-            # with gr.Row():
-            model_name = gr.Dropdown(
-                label="Model",
-                scale=3,
-                choices=["detr_demo_boxes", "detr_resnet101_panoptic"],
-                value="detr_demo_boxes",
-            )
+def export_model(model_name, progress=gr.Progress(track_tqdm=True)):
+    progress(0.1, "Conversion..")
+    t0 = perf_counter()
+    if model_name == "detr_simple_demo":
+        model = SimpleDetr()
+    elif model_name == "detr_resnet101_panoptic":
+        model = PanopticDetrResenet101()
+
+    model_path = model.export()
+    t1 = perf_counter()
+    return model_path, {"export_time": t1 - t0}
+
 
-            conf = gr.Slider(label="Confidence", minimum=0, maximum=0.99, value=0.5)
+with gr.Blocks() as demo:
+    gr.Markdown("# DETR: Detection Transformer")
+    # gr.Image(value=f"{ASSETS_DIR}/detr_architecture.png")
+    with gr.Tab("Torch Inference"):
+        with gr.Row():
+            with gr.Column():
+                img_file = gr.Image(type="pil")
+                model_name = gr.Dropdown(
+                    label="Model",
+                    choices=[
+                        "detr_simple_demo",
+                        "detr_resnet101_panoptic",
+                    ],
+                    value="detr_simple_demo",
+                )
 
-            with gr.Row():
-                start_btn = gr.Button("Start", variant="primary")
-        with gr.Column():
-            annotated_img = gr.Image(label="Annotated Image")
-            speed = gr.JSON(label="speed")
-    examples = gr.Examples(
-        examples=[
-            [path]
-            for path in sv.list_files_with_extensions(
-                directory=ASSETS_DIR, extensions=["jpeg", "jpg", "png"]
-            )
-        ],
-        inputs=[img_file],
-    )
-    start_btn.click(
-        fn=run_inference,
-        inputs=[img_file, conf, model_name],
-        outputs=[annotated_img, speed],
-    )
+                conf = gr.Slider(label="Confidence", minimum=0, maximum=0.99, value=0.5)
+
+                with gr.Row():
+                    start_btn = gr.Button("Start", variant="primary")
+            with gr.Column():
+                annotated_img = gr.Image(label="Annotated Image")
+                speed = gr.JSON(label="speed")
+        examples = gr.Examples(
+            examples=[
+                [path]
+                for path in sv.list_files_with_extensions(
+                    directory=IMAGES_DIR, extensions=["jpeg", "jpg", "png"]
+                )
+            ],
+            inputs=[img_file],
+        )
+        start_btn.click(
+            fn=run_inference,
+            inputs=[img_file, conf, model_name],
+            outputs=[annotated_img, speed],
+        )
+    with gr.Tab("ONNX Inference"):
+        with gr.Row():
+            with gr.Column():
+                img_file = gr.Image(type="pil")
+                model_name = gr.Dropdown(
+                    label="Model",
+                    choices=[
+                        "detr_simple_demo_onnx",
+                    ],
+                    value="detr_simple_demo_onnx",
+                )
+                conf = gr.Slider(label="Confidence", minimum=0, maximum=0.99, value=0.7)
+                with gr.Row():
+                    start_btn = gr.Button("Start", variant="primary")
+            with gr.Column():
+                annotated_img = gr.Image(label="Annotated Image")
+                speed = gr.JSON(label="speed")
+        examples = gr.Examples(
+            examples=[
+                [path]
+                for path in sv.list_files_with_extensions(
+                    directory=IMAGES_DIR, extensions=["jpeg", "jpg", "png"]
+                )
+            ],
+            inputs=[img_file],
+        )
+        start_btn.click(
+            fn=run_inference,
+            inputs=[img_file, conf, model_name],
+            outputs=[annotated_img, speed],
+        )
+    with gr.Tab("ONNX export"):
+        with gr.Row():
+            with gr.Column():
+                model_name = gr.Dropdown(
+                    label="Model",
+                    choices=[
+                        "detr_simple_demo",
+                        "detr_resnet101_panoptic",
+                    ],
+                    value="detr_simple_demo",
+                )
+                with gr.Row():
+                    export_btn = gr.Button("Export", variant="primary")
+            with gr.Column():
+                onnx_file = gr.File()
+                result = gr.JSON(label="result")
+        export_btn.click(
+            fn=export_model,
+            inputs=[model_name],
+            outputs=[onnx_file, result],
+        )
 
 if __name__ == "__main__":
-    inference_gradio.queue(2).launch(
+    demo.queue(2).launch(
         debug=True,
         server_name="0.0.0.0",
         server_port=7000,

requirements.txt

@@ -5,4 +5,6 @@ matplotlib
 torchvision
 supervision==0.17.1
 git+https://github.com/cocodataset/panopticapi.git
-seaborn
+seaborn
+onnx
+onnxruntime