From d69e3114ae87a9632e25262a3b452b52d729cd0e Mon Sep 17 00:00:00 2001 From: zhiqwang Date: Fri, 9 Apr 2021 03:33:26 -0400 Subject: [PATCH] Cleanup unittest --- test/test_image_utils.py | 24 ++--------------- test/test_models.py | 5 +--- test/test_onnx.py | 58 +++++++++++++++++++--------------------- test/test_utils.py | 23 ++++++++++++++-- 4 files changed, 51 insertions(+), 59 deletions(-) diff --git a/test/test_image_utils.py b/test/test_image_utils.py index 8fc10e10..16d7d62d 100644 --- a/test/test_image_utils.py +++ b/test/test_image_utils.py @@ -1,33 +1,13 @@ +# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved. import unittest import numpy as np import torch -from torch import Tensor -from yolort.utils.image_utils import ( - box_cxcywh_to_xyxy, - letterbox, - read_image_to_tensor, - get_image_from_url, - scale_coords, -) +from yolort.utils.image_utils import box_cxcywh_to_xyxy, letterbox, scale_coords class ImageUtilsTester(unittest.TestCase): - def test_read_image(self): - N, H, W = 3, 720, 360 - img = np.random.randint(0, 255, (H, W, N), dtype='uint8') # As a dummy image - out = read_image_to_tensor(img) - - self.assertIsInstance(out, Tensor) - self.assertEqual(tuple(out.shape), (N, H, W)) - - def test_get_image_from_url(self): - url = "https://gitee.com/zhiqwang/yolov5-rt-stack/raw/master/test/assets/zidane.jpg" - img = get_image_from_url(url) - self.assertIsInstance(img, np.ndarray) - self.assertTupleEqual(img.shape, (720, 1280, 3)) - def test_letterbox(self): img = np.random.randint(0, 255, (720, 360, 3), dtype='uint8') # As a dummy image out = letterbox(img, new_shape=(416, 416))[0] diff --git a/test/test_models.py b/test/test_models.py index f4ba0602..57451fd1 100644 --- a/test/test_models.py +++ b/test/test_models.py @@ -1,3 +1,4 @@ +# Copyright (c) 2020, Zhiqiang Wang. All Rights Reserved. import unittest import torch @@ -239,7 +240,3 @@ def test_anchor_generator(self): self.assertEqual(anchors[0], anchor_output) self.assertEqual(anchors[1], wh_output) self.assertEqual(anchors[2], xy_output) - - -if __name__ == '__main__': - unittest.main() diff --git a/test/test_onnx.py b/test/test_onnx.py index 4280eb23..c89328ac 100644 --- a/test/test_onnx.py +++ b/test/test_onnx.py @@ -1,16 +1,22 @@ +# Copyright (c) 2020, Zhiqiang Wang. All Rights Reserved. +""" +Test for exporting model to ONNX and inference with ONNXRuntime +""" import io -import torch +import unittest -# onnxruntime requires python 3.5 or above try: + # This import should be before that of torch if you are using PyTorch lower than 1.5.0 + # see import onnxruntime except ImportError: onnxruntime = None -import unittest +import torch from torchvision.ops._register_onnx_ops import _onnx_opset_version from yolort.models import yolov5s, yolov5m, yolotr +from yolort.utils import get_image_from_url, read_image_to_tensor @unittest.skipIf(onnxruntime is None, 'ONNX Runtime unavailable') @@ -22,13 +28,21 @@ def setUpClass(cls): def run_model(self, model, inputs_list, tolerate_small_mismatch=False, do_constant_folding=True, dynamic_axes=None, output_names=None, input_names=None): + """ + The core part of exporting model to ONNX and inference with ONNXRuntime + Copy-paste from + """ model.eval() onnx_io = io.BytesIO() + if isinstance(inputs_list[0][-1], dict): + torch_onnx_input = inputs_list[0] + ({},) + else: + torch_onnx_input = inputs_list[0] # export to onnx with the first input torch.onnx.export( model, - inputs_list[0], + torch_onnx_input, onnx_io, do_constant_folding=do_constant_folding, opset_version=_onnx_opset_version, @@ -74,28 +88,14 @@ def to_numpy(tensor): else: raise - def get_image_from_url(self, url, size=None): - import requests - from PIL import Image - from io import BytesIO - from torchvision import transforms - - data = requests.get(url) - image = Image.open(BytesIO(data.content)).convert("RGB") - - if size is None: - size = (300, 200) - image = image.resize(size, Image.BILINEAR) - - to_tensor = transforms.ToTensor() - return to_tensor(image) - def get_test_images(self): - image_url = "http://farm3.staticflickr.com/2469/3915380994_2e611b1779_z.jpg" - image = self.get_image_from_url(url=image_url, size=(100, 320)) + image_url = "https://github.com/ultralytics/yolov5/raw/master/data/images/bus.jpg" + image = get_image_from_url(image_url) + image = read_image_to_tensor(image, is_half=False) - image_url2 = "https://pytorch.org/tutorials/_static/img/tv_tutorial/tv_image05.png" - image2 = self.get_image_from_url(url=image_url2, size=(250, 380)) + image_url2 = "https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg" + image2 = get_image_from_url(image_url2) + image2 = read_image_to_tensor(image2, is_half=False) images_one = [image] images_two = [image2] @@ -104,7 +104,7 @@ def get_test_images(self): def test_yolov5s_r31(self): images_one, images_two = self.get_test_images() images_dummy = [torch.ones(3, 100, 100) * 0.3] - model = yolov5s(upstream_version='r3.1', export_friendly=True, pretrained=True) + model = yolov5s(upstream_version='r3.1', export_friendly=True, pretrained=True, score_thresh=0.45) model.eval() model(images_one) # Test exported model on images of different size, or dummy input @@ -121,7 +121,7 @@ def test_yolov5s_r31(self): def test_yolov5m_r40(self): images_one, images_two = self.get_test_images() images_dummy = [torch.ones(3, 100, 100) * 0.3] - model = yolov5m(upstream_version='r4.0', export_friendly=True, pretrained=True) + model = yolov5m(upstream_version='r4.0', export_friendly=True, pretrained=True, score_thresh=0.45) model.eval() model(images_one) # Test exported model on images of different size, or dummy input @@ -138,7 +138,7 @@ def test_yolov5m_r40(self): def test_yolotr(self): images_one, images_two = self.get_test_images() images_dummy = [torch.ones(3, 100, 100) * 0.3] - model = yolotr(upstream_version='r4.0', export_friendly=True, pretrained=True) + model = yolotr(upstream_version='r4.0', export_friendly=True, pretrained=True, score_thresh=0.45) model.eval() model(images_one) # Test exported model on images of different size, or dummy input @@ -151,7 +151,3 @@ def test_yolotr(self): output_names=["outputs"], dynamic_axes={"images_tensors": [0, 1, 2], "outputs": [0, 1, 2]}, tolerate_small_mismatch=True) - - -if __name__ == '__main__': - unittest.main() diff --git a/test/test_utils.py b/test/test_utils.py index e3808489..766a482e 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -1,9 +1,14 @@ # Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved. import unittest +import numpy as np -from torch import nn +from torch import nn, Tensor -from yolort.utils import update_module_state_from_ultralytics +from yolort.utils import ( + update_module_state_from_ultralytics, + read_image_to_tensor, + get_image_from_url, +) class UtilsTester(unittest.TestCase): @@ -16,3 +21,17 @@ def test_update_module_state_from_ultralytics(self): custom_path_or_model=None, ) self.assertIsInstance(model, nn.Module) + + def test_read_image_to_tensor(self): + N, H, W = 3, 720, 360 + img = np.random.randint(0, 255, (H, W, N), dtype='uint8') # As a dummy image + out = read_image_to_tensor(img) + + self.assertIsInstance(out, Tensor) + self.assertEqual(tuple(out.shape), (N, H, W)) + + def test_get_image_from_url(self): + url = "https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg" + img = get_image_from_url(url) + self.assertIsInstance(img, np.ndarray) + self.assertTupleEqual(img.shape, (720, 1280, 3))