From 76b4d0577d0dae4e366cb8fa4a8346864986a59a Mon Sep 17 00:00:00 2001
From: George Hotz <george@comma.ai>
Date: Wed, 22 Feb 2023 19:32:01 -0800
Subject: [PATCH] yolov8 works up to the MaxPool

---
 examples/yolov8.py | 17 +++++++++++++++++
 extra/onnx.py      |  6 ++++--
 2 files changed, 21 insertions(+), 2 deletions(-)
 create mode 100644 examples/yolov8.py

diff --git a/examples/yolov8.py b/examples/yolov8.py
new file mode 100644
index 0000000000..53c2c0a39c
--- /dev/null
+++ b/examples/yolov8.py
@@ -0,0 +1,17 @@
+#!/usr/bin/env python3
+import os
+from ultralytics import YOLO
+import onnx
+from extra.onnx import get_run_onnx
+from tinygrad.tensor import Tensor
+
+os.chdir("/tmp")
+if not os.path.isfile("yolov8n-seg.onnx"):
+  model = YOLO("yolov8n-seg.pt") 
+  model.export(format="onnx", imgsz=[480,640])
+onnx_model = onnx.load(open("yolov8n-seg.onnx", "rb"))
+# TODO: move get example inputs to onnx
+input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
+print(input_shapes)
+run_onnx = get_run_onnx(onnx_model)
+run_onnx({"images": Tensor.zeros(1,3,480,640)}, debug=True)
diff --git a/extra/onnx.py b/extra/onnx.py
index 7af4e7c8ea..65fc2a7214 100644
--- a/extra/onnx.py
+++ b/extra/onnx.py
@@ -18,6 +18,7 @@ def get_run_onnx(onnx_model):
   def attribute_parse(a):
     if a.type == 7: return tuple([int(x) for x in a.ints])
     elif a.type == 4: return buffer_parse(a.t)  # TENSOR
+    elif a.type == 3: return str(a.s)
     elif a.type == 2: return int(a.i)
     elif a.type == 1: return float(a.f)
     else: raise Exception(f"can't parse {a.type} {a}")
@@ -136,6 +137,7 @@ def get_run_onnx(onnx_model):
         if n.op_type == "Sub": ret = inp[0] - inp[1]
         if n.op_type == "Mul": ret = inp[0] * inp[1]
       elif n.op_type == "Split":
+        if 'split' not in opt: opt['split'] = [int(x) for x in inp[1].numpy()]  # split can be a tensor
         i = 0
         arg = [(0,x) for x in inp[0].shape]
         for o,s in zip(n.output, opt['split']):
@@ -147,11 +149,11 @@ def get_run_onnx(onnx_model):
         assert opt['kernel_shape'] == opt['strides'] or opt['strides'] == (1,1)
         ret = inp[0].avg_pool2d(opt['kernel_shape'])
       elif n.op_type == "MaxPool":
-        assert opt['kernel_shape'] == opt['strides']
+        assert opt['kernel_shape'] == opt['strides'], f"kernel_shape and stride mismatch {opt}"
         #opt['kernel_shape'] = opt['strides']
         # TODO: this is untested and probably wrong
         ret = inp[0].pad2d(opt['pads'])
-        ret = ret.max_pool2d(opt['kernel_shape'])
+        ret = ret.max_pool2d(opt['kernel_shape'], opt['strides'])
         # strides aren't supported in max_pool
         #chan = ret.shape[1]
         #w = Tensor.eye(chan).reshape((chan, chan, 1, 1))