EasyCV/tests/predictors/test_detector_blade.py
zouxinyi0625 9ad1a8f448
Fix YOLOX-Edge problem and update yolox-pai turtorial (#182)
* complete yolox introduction

* fix env problem with nni2.6.1

* fix yolox-edge problem

*  add docker turtorial

* add blade test

* support old export yolox model

Co-authored-by: wuziheng <zhoulou.wzh@alibaba-inc.com>
2022-09-06 17:11:12 +08:00

117 lines
4.4 KiB
Python

# Copyright (c) Alibaba, Inc. and its affiliates.
"""
isort:skip_file
"""
import os
import tempfile
import unittest
import cv2
import numpy as np
from PIL import Image
from easycv.predictors.detector import TorchYoloXPredictor
from tests.ut_config import (PRETRAINED_MODEL_YOLOXS_EXPORT,
PRETRAINED_MODEL_YOLOXS_NOPRE_NOTRT_JIT,
PRETRAINED_MODEL_YOLOXS_PRE_NOTRT_JIT,
PRETRAINED_MODEL_YOLOXS_NOPRE_NOTRT_BLADE,
PRETRAINED_MODEL_YOLOXS_PRE_NOTRT_BLADE,
DET_DATA_SMALL_COCO_LOCAL)
from easycv.utils.test_util import benchmark
import logging
import pandas as pd
import torch
from numpy.testing import assert_array_almost_equal
@unittest.skipIf(torch.__version__ != '1.8.1+cu102',
'Blade need another environment')
class DetectorTest(unittest.TestCase):
def setUp(self):
print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
def test_yolox_detector_blade_nopre_notrt(self):
img = os.path.join(DET_DATA_SMALL_COCO_LOCAL,
'val2017/000000522713.jpg')
input_data_list = [np.asarray(Image.open(img))]
blade_path = PRETRAINED_MODEL_YOLOXS_NOPRE_NOTRT_BLADE
# blade_path = '/home/zouxinyi.zxy/easycv_nfs/pretrained_models/detection/infer_yolox/debug_blade.pt.blade'
predictor_blade = TorchYoloXPredictor(
model_path=blade_path, score_thresh=0.5)
output = predictor_blade.predict(input_data_list)[0]
self.assertIn('detection_boxes', output)
self.assertIn('detection_scores', output)
self.assertIn('detection_classes', output)
self.assertIn('detection_class_names', output)
self.assertIn('ori_img_shape', output)
self.assertEqual(len(output['detection_boxes']), 4)
self.assertEqual(output['ori_img_shape'], [480, 640])
self.assertListEqual(output['detection_classes'].tolist(),
np.array([13, 8, 8, 8], dtype=np.int32).tolist())
self.assertListEqual(output['detection_class_names'],
['bench', 'boat', 'boat', 'boat'])
assert_array_almost_equal(
output['detection_scores'],
np.array([0.92593855, 0.60268813, 0.57775956, 0.5750004],
dtype=np.float32),
decimal=2)
assert_array_almost_equal(
output['detection_boxes'],
np.array([[407.89523, 284.62598, 561.4984, 356.7296],
[439.37653, 263.42395, 467.01526, 271.79144],
[480.8597, 269.64435, 502.18765, 274.80127],
[510.37033, 268.4982, 527.67017, 273.04935]]),
decimal=1)
def test_yolox_detector_blade_pre_notrt(self):
img = os.path.join(DET_DATA_SMALL_COCO_LOCAL,
'val2017/000000522713.jpg')
input_data_list = [np.asarray(Image.open(img))]
blade_path = PRETRAINED_MODEL_YOLOXS_PRE_NOTRT_BLADE
predictor_blade = TorchYoloXPredictor(
model_path=blade_path, score_thresh=0.5)
output = predictor_blade.predict(input_data_list)[0]
self.assertIn('detection_boxes', output)
self.assertIn('detection_scores', output)
self.assertIn('detection_classes', output)
self.assertIn('detection_class_names', output)
self.assertIn('ori_img_shape', output)
self.assertEqual(len(output['detection_boxes']), 4)
self.assertEqual(output['ori_img_shape'], [480, 640])
self.assertListEqual(output['detection_classes'].tolist(),
np.array([13, 8, 8, 8], dtype=np.int32).tolist())
self.assertListEqual(output['detection_class_names'],
['bench', 'boat', 'boat', 'boat'])
assert_array_almost_equal(
output['detection_scores'],
np.array([0.92593855, 0.60268813, 0.57775956, 0.5750004],
dtype=np.float32),
decimal=2)
assert_array_almost_equal(
output['detection_boxes'],
np.array([[407.89523, 284.62598, 561.4984, 356.7296],
[439.37653, 263.42395, 467.01526, 271.79144],
[480.8597, 269.64435, 502.18765, 274.80127],
[510.37033, 268.4982, 527.67017, 273.04935]]),
decimal=1)
if __name__ == '__main__':
unittest.main()