You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
69 lines
3.6 KiB
69 lines
3.6 KiB
import numpy as np
|
|
from examples.yolov8 import YOLOv8, get_variant_multiples, preprocess, postprocess, label_predictions
|
|
import unittest
|
|
import io, cv2
|
|
import onnxruntime as ort
|
|
import ultralytics
|
|
from tinygrad.nn.state import safe_load, load_state_dict
|
|
from tinygrad.helpers import fetch
|
|
|
|
class TestYOLOv8(unittest.TestCase):
|
|
def test_all_load_weights(self):
|
|
for variant in ['n', 's', 'm', 'l', 'x']:
|
|
depth, width, ratio = get_variant_multiples(variant)
|
|
TinyYolov8 = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
|
|
state_dict = safe_load(fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{variant}.safetensors'))
|
|
load_state_dict(TinyYolov8, state_dict)
|
|
print(f'successfully loaded weights for yolov{variant}')
|
|
|
|
def test_predictions(self):
|
|
test_image_urls = ['https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/bus.jpg', 'https://www.aljazeera.com/wp-content/uploads/2022/10/2022-04-28T192650Z_1186456067_UP1EI4S1I0P14_RTRMADP_3_SOCCER-ENGLAND-MUN-CHE-REPORT.jpg']
|
|
variant = 'n'
|
|
depth, width, ratio = get_variant_multiples(variant)
|
|
TinyYolov8 = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
|
|
state_dict = safe_load(fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{variant}.safetensors'))
|
|
load_state_dict(TinyYolov8, state_dict)
|
|
|
|
for i in range(len(test_image_urls)):
|
|
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}
|
|
|
|
def test_forward_pass_torch_onnx(self):
|
|
variant = 'n'
|
|
weights_location = fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{variant}.safetensors')
|
|
weights_location_pt = fetch(f'https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8{variant}.pt', name=f"yolov8{variant}.pt") # it needs the pt extension # noqa: E501
|
|
weights_location_onnx = weights_location_pt.parent / f"yolov8{variant}.onnx"
|
|
|
|
# the ultralytics export prints a lot of unneccesary things
|
|
if not weights_location_onnx.is_file():
|
|
model = ultralytics.YOLO(model=weights_location_pt, task='Detect')
|
|
model.export(format="onnx",imgsz=[640, 480])
|
|
|
|
depth, width, ratio = get_variant_multiples(variant)
|
|
TinyYolov8 = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
|
|
state_dict = safe_load(weights_location)
|
|
load_state_dict(TinyYolov8, state_dict)
|
|
|
|
image_location = [np.frombuffer(io.BytesIO(fetch('https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/bus.jpg').read_bytes()).read(), np.uint8)] # noqa: E501
|
|
orig_image = [cv2.imdecode(image_location[0], 1)]
|
|
|
|
input_image = preprocess(orig_image)
|
|
|
|
onnx_session = ort.InferenceSession(weights_location_onnx)
|
|
onnx_input_name = onnx_session.get_inputs()[0].name
|
|
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)
|
|
|
|
# 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)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|
|
|