More yolo processing in tinygrad (#9928)

* more tg less np

* update webgpu html for new compile

* resize boxes

* remove text

* add back note

* fix indentation

* fix indentation

* remove magic num

* remove now unused funcs

* back to numpy nms

* no loop

* fix iou suppression

* update test

* dont suppress other classes

* add working scale

* fix expected value, rounded up 0.24 was being counted

* add postprocess bool for onnx test

* fix indents

* clean

* clean

* fix indent

* remove print

* fix indent

* remove unused import

* remove hardcoded 0.25

* space

* spacing

* clean label_predictions func

* remove single item lists

* space

* use postprocess output in test

* space

* clean

* clean

* remove redundant threshold

* remove redundant threshold

* clean

* rename var

* move loop into func

* unhardcode iou_threshold

* remove unused values

* clean

* add note

* clean

* keep const

* move back funcs

---------

Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com>
This commit is contained in:
Rory Clear
2025-04-24 21:21:46 +01:00
committed by GitHub
parent 74c6cf8be3
commit 3a189fa561
3 changed files with 98 additions and 208 deletions

View File

@@ -1,5 +1,6 @@
import numpy as np
from examples.yolov8 import YOLOv8, get_variant_multiples, preprocess, postprocess, label_predictions
from examples.yolov8 import YOLOv8, get_variant_multiples, preprocess, label_predictions, postprocess
from tinygrad import Tensor
import unittest
import io, cv2
import onnxruntime as ort
@@ -28,9 +29,8 @@ class TestYOLOv8(unittest.TestCase):
img = cv2.imdecode(np.frombuffer(fetch(test_image_urls[i]).read_bytes(), np.uint8), 1)
test_image = preprocess([img])
predictions = TinyYolov8(test_image)
post_predictions = postprocess(preds=predictions, img=test_image, orig_imgs=[img])
labels = label_predictions(post_predictions)
assert labels == {5: 1, 0: 4, 11: 1} if i == 0 else labels == {0: 13, 29: 1, 32: 1}
labels = label_predictions(predictions.numpy())
assert labels == {5: 1, 0: 4, 11: 1} if i == 0 else labels == {0: 12, 29: 1, 32: 1}
def test_forward_pass_torch_onnx(self):
variant = 'n'
@@ -58,12 +58,16 @@ class TestYOLOv8(unittest.TestCase):
onnx_output_name = onnx_session.get_outputs()[0].name
onnx_output = onnx_session.run([onnx_output_name], {onnx_input_name: input_image.numpy()})
tiny_output = TinyYolov8(input_image)
tiny_output = TinyYolov8(input_image).numpy()
onnx_output = postprocess(Tensor(onnx_output[0])).numpy()
#invalid boxes are multiplied by zero in postprocess
onnx_output = onnx_output[onnx_output[:, 4] != 0]
tiny_output = tiny_output[tiny_output[:, 4] != 0]
# currently rtol is 0.025 because there is a 1-2% difference in our predictions
# because of the zero padding in SPPF module (line 280) maxpooling layers rather than the -infinity in torch.
# This difference does not make a difference "visually".
np.testing.assert_allclose(onnx_output[0], tiny_output.numpy(), atol=5e-4, rtol=0.025)
np.testing.assert_allclose(onnx_output, tiny_output, atol=5e-4, rtol=0.025)
if __name__ == '__main__':
unittest.main()