mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-12 07:23:45 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
74 lines
3.8 KiB
Python
74 lines
3.8 KiB
Python
import numpy as np
|
|
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
|
|
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)
|
|
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'
|
|
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).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, tiny_output, atol=5e-4, rtol=0.01)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|