Use test images from repo rather than internet (#149)

* Fixing lint

* Fixing unittest

* Use test images from repo rather than internet

* Fixing unittest

test/test_onnx.py

@@ -2,25 +2,26 @@
 """
 Test for exporting model to ONNX and inference with ONNXRuntime
 """
-import io
-import unittest
+from typing import List, Tuple
 
-try:
-    # This import should be before that of torch if you are using PyTorch lower than 1.5.0
-    # see <https://github.com/onnx/onnx/issues/2394#issuecomment-581638840>
-    import onnxruntime
-except ImportError:
-    onnxruntime = None
+from pathlib import Path
+import io
+import pytest
+from PIL import Image
 
 import torch
+from torch import Tensor
+from torchvision import transforms
 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
 
+# In environments without onnxruntime we prefer to
+# invoke all tests in the repo and have this one skipped rather than fail.
+onnxruntime = pytest.importorskip("onnxruntime")
 
-@unittest.skipIf(onnxruntime is None, 'ONNX Runtime unavailable')
-class ONNXExporterTester(unittest.TestCase):
+
+class TestONNXExporter:
     @classmethod
     def setUpClass(cls):
         torch.manual_seed(123)
@@ -53,10 +54,10 @@ def run_model(self, model, inputs_list, tolerate_small_mismatch=False,
         # validate the exported model with onnx runtime
         for test_inputs in inputs_list:
             with torch.no_grad():
-                if isinstance(test_inputs, torch.Tensor) or isinstance(test_inputs, list):
+                if isinstance(test_inputs, Tensor) or isinstance(test_inputs, list):
                     test_inputs = (test_inputs,)
                 test_ouputs = model(*test_inputs)
-                if isinstance(test_ouputs, torch.Tensor):
+                if isinstance(test_ouputs, Tensor):
                     test_ouputs = (test_ouputs,)
             self.ort_validate(onnx_io, test_inputs, test_ouputs, tolerate_small_mismatch)
 
@@ -88,18 +89,16 @@ def to_numpy(tensor):
                 else:
                     raise
 
-    def get_test_images(self):
-        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)
+    def get_image(self, img_name, size) -> Tensor:
+
+        img_path = Path(__file__).parent.resolve() / "assets" / img_name
+        image = Image.open(img_path).convert("RGB").resize(size, Image.BILINEAR)
 
-        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)
+        return transforms.ToTensor()(image)
 
-        images_one = [image]
-        images_two = [image2]
-        return images_one, images_two
+    def get_test_images(self) -> Tuple[List[Tensor], List[Tensor]]:
+        return ([self.get_image("bus.jpg", (416, 320))],
+                [self.get_image("zidane.jpg", (352, 480))])
 
     def test_yolov5s_r31(self):
         images_one, images_two = self.get_test_images()

test/test_torchscript.py

@@ -1,68 +1,69 @@
 # Copyright (c) 2020, Zhiqiang Wang. All Rights Reserved.
-import unittest
-
 import torch
 
 from yolort.models import yolov5s, yolov5m, yolov5l, yolotr
 
 
-class TorchScriptTester(unittest.TestCase):
-    def test_yolov5s_script(self):
-        model = yolov5s(pretrained=True)
-        model.eval()
+def test_yolov5s_script():
+    model = yolov5s(pretrained=True, size=(320, 320), score_thresh=0.45)
+    model.eval()
+
+    scripted_model = torch.jit.script(model)
+    scripted_model.eval()
+
+    x = [torch.rand(3, 288, 320), torch.rand(3, 300, 256)]
+
+    out = model(x)
+    out_script = scripted_model(x)
+
+    torch.testing.assert_allclose(out[0]["scores"], out_script[1][0]["scores"], rtol=0., atol=0.)
+    torch.testing.assert_allclose(out[0]["labels"], out_script[1][0]["labels"], rtol=0., atol=0.)
+    torch.testing.assert_allclose(out[0]["boxes"], out_script[1][0]["boxes"], rtol=0., atol=0.)
 
-        scripted_model = torch.jit.script(model)
-        scripted_model.eval()
 
-        x = [torch.rand(3, 416, 320), torch.rand(3, 480, 352)]
+def test_yolov5m_script():
+    model = yolov5m(pretrained=True, size=(320, 320), score_thresh=0.45)
+    model.eval()
 
-        out = model(x)
-        out_script = scripted_model(x)
-        self.assertTrue(out[0]["scores"].equal(out_script[1][0]["scores"]))
-        self.assertTrue(out[0]["labels"].equal(out_script[1][0]["labels"]))
-        self.assertTrue(out[0]["boxes"].equal(out_script[1][0]["boxes"]))
+    scripted_model = torch.jit.script(model)
+    scripted_model.eval()
 
-    def test_yolov5m_script(self):
-        model = yolov5m(pretrained=True)
-        model.eval()
+    x = [torch.rand(3, 288, 320), torch.rand(3, 300, 256)]
 
-        scripted_model = torch.jit.script(model)
-        scripted_model.eval()
+    out = model(x)
+    out_script = scripted_model(x)
+    torch.testing.assert_allclose(out[0]["scores"], out_script[1][0]["scores"], rtol=0., atol=0.)
+    torch.testing.assert_allclose(out[0]["labels"], out_script[1][0]["labels"], rtol=0., atol=0.)
+    torch.testing.assert_allclose(out[0]["boxes"], out_script[1][0]["boxes"], rtol=0., atol=0.)
 
-        x = [torch.rand(3, 416, 320), torch.rand(3, 480, 352)]
 
-        out = model(x)
-        out_script = scripted_model(x)
-        self.assertTrue(out[0]["scores"].equal(out_script[1][0]["scores"]))
-        self.assertTrue(out[0]["labels"].equal(out_script[1][0]["labels"]))
-        self.assertTrue(out[0]["boxes"].equal(out_script[1][0]["boxes"]))
+def test_yolov5l_script():
+    model = yolov5l(pretrained=True, size=(320, 320), score_thresh=0.45)
+    model.eval()
 
-    def test_yolov5l_script(self):
-        model = yolov5l(pretrained=True)
-        model.eval()
+    scripted_model = torch.jit.script(model)
+    scripted_model.eval()
 
-        scripted_model = torch.jit.script(model)
-        scripted_model.eval()
+    x = [torch.rand(3, 288, 320), torch.rand(3, 300, 256)]
 
-        x = [torch.rand(3, 416, 320), torch.rand(3, 480, 352)]
+    out = model(x)
+    out_script = scripted_model(x)
+    torch.testing.assert_allclose(out[0]["scores"], out_script[1][0]["scores"], rtol=0., atol=0.)
+    torch.testing.assert_allclose(out[0]["labels"], out_script[1][0]["labels"], rtol=0., atol=0.)
+    torch.testing.assert_allclose(out[0]["boxes"], out_script[1][0]["boxes"], rtol=0., atol=0.)
 
-        out = model(x)
-        out_script = scripted_model(x)
-        self.assertTrue(out[0]["scores"].equal(out_script[1][0]["scores"]))
-        self.assertTrue(out[0]["labels"].equal(out_script[1][0]["labels"]))
-        self.assertTrue(out[0]["boxes"].equal(out_script[1][0]["boxes"]))
 
-    def test_yolotr_script(self):
-        model = yolotr(pretrained=True)
-        model.eval()
+def test_yolotr_script():
+    model = yolotr(pretrained=True, size=(320, 320), score_thresh=0.45)
+    model.eval()
 
-        scripted_model = torch.jit.script(model)
-        scripted_model.eval()
+    scripted_model = torch.jit.script(model)
+    scripted_model.eval()
 
-        x = [torch.rand(3, 416, 320), torch.rand(3, 480, 352)]
+    x = [torch.rand(3, 288, 320), torch.rand(3, 300, 256)]
 
-        out = model(x)
-        out_script = scripted_model(x)
-        self.assertTrue(out[0]["scores"].equal(out_script[1][0]["scores"]))
-        self.assertTrue(out[0]["labels"].equal(out_script[1][0]["labels"]))
-        self.assertTrue(out[0]["boxes"].equal(out_script[1][0]["boxes"]))
+    out = model(x)
+    out_script = scripted_model(x)
+    torch.testing.assert_allclose(out[0]["scores"], out_script[1][0]["scores"], rtol=0., atol=0.)
+    torch.testing.assert_allclose(out[0]["labels"], out_script[1][0]["labels"], rtol=0., atol=0.)
+    torch.testing.assert_allclose(out[0]["boxes"], out_script[1][0]["boxes"], rtol=0., atol=0.)

yolort/utils/image_utils.py

@@ -145,8 +145,8 @@ def scale_coords(coords, img_shape, img_shape_origin, ratio_pad=None):
     return coords
 
 
-def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
-                        labels=(), max_det=300):
+def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None,
+                        agnostic=False, multi_label=False, labels=(), max_det=300):
     """Runs Non-Maximum Suppression (NMS) on inference results
 
     Returns:
@@ -157,11 +157,13 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
     xc = prediction[..., 4] > conf_thres  # candidates
 
     # Checks
-    assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
-    assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
+    assert 0 <= conf_thres <= 1, (
+        f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0')
+    assert 0 <= iou_thres <= 1, (
+        f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0')
 
     # Settings
-    min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
+    max_wh = 4096  # (pixels) minimum and maximum box width and height
     max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
     time_limit = 10.0  # seconds to quit after
     redundant = True  # require redundant detections
@@ -177,11 +179,11 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
 
         # Cat apriori labels if autolabelling
         if labels and len(labels[xi]):
-            l = labels[xi]
-            v = torch.zeros((len(l), nc + 5), device=x.device)
-            v[:, :4] = l[:, 1:5]  # box
+            label = labels[xi]
+            v = torch.zeros((len(label), nc + 5), device=x.device)
+            v[:, :4] = label[:, 1:5]  # box
             v[:, 4] = 1.0  # conf
-            v[range(len(l)), l[:, 0].long() + 5] = 1.0  # cls
+            v[range(len(label)), label[:, 0].long() + 5] = 1.0  # cls
             x = torch.cat((x, v), 0)
 
         # If none remain process next image
@@ -227,7 +229,8 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
             # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
             iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
             weights = iou * scores[None]  # box weights
-            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
+            # merging boxes
+            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)
             if redundant:
                 i = i[iou.sum(1) > 1]  # require redundancy
 
@@ -238,6 +241,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
 
     return output
 
+
 def get_image_from_url(
     url: str,
     flags: int = cv2.IMREAD_COLOR,