|
- import argparse
- import os.path as osp
- import warnings
- import numbers
- from functools import partial
- import sys
- sys.path.append("/home/shanwei-luo/userdata/mmdetection")
- import numpy as np
- #import onnx
- import torch
- import cv2
- from mmcv import Config, DictAction
-
- from mmdet.core.export.model_wrappers import ONNXRuntimeDetector
-
- from mmdet.apis import (async_inference_detector, inference_detector,
- init_detector, show_result_pyplot)
-
- import onnxruntime as ort
- import onnx
- import mmcv
-
- def impad(img,
- *,
- shape=None,
- padding=None,
- pad_val=0,
- padding_mode='constant'):
- assert (shape is not None) ^ (padding is not None)
- if shape is not None:
- padding = (0, 0, shape[1] - img.shape[1], shape[0] - img.shape[0])
-
- # check pad_val
- if isinstance(pad_val, tuple):
- assert len(pad_val) == img.shape[-1]
- elif not isinstance(pad_val, numbers.Number):
- raise TypeError('pad_val must be a int or a tuple. '
- f'But received {type(pad_val)}')
-
- # check padding
- if isinstance(padding, tuple) and len(padding) in [2, 4]:
- if len(padding) == 2:
- padding = (padding[0], padding[1], padding[0], padding[1])
- elif isinstance(padding, numbers.Number):
- padding = (padding, padding, padding, padding)
- else:
- raise ValueError('Padding must be a int or a 2, or 4 element tuple.'
- f'But received {padding}')
-
- # check padding mode
- assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric']
-
- border_type = {
- 'constant': cv2.BORDER_CONSTANT,
- 'edge': cv2.BORDER_REPLICATE,
- 'reflect': cv2.BORDER_REFLECT_101,
- 'symmetric': cv2.BORDER_REFLECT
- }
- img = cv2.copyMakeBorder(
- img,
- padding[1],
- padding[3],
- padding[0],
- padding[2],
- border_type[padding_mode],
- value=pad_val)
-
- return img
-
- def imnormalize(img, mean, std, to_rgb=True):
- mean = np.float64(mean.reshape(1, -1))
- stdinv = 1 / np.float64(std.reshape(1, -1))
- if to_rgb:
- cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img) # inplace
- cv2.subtract(img, mean, img) # inplace
- cv2.multiply(img, stdinv, img) # inplace
- return img
-
- def preprocess_example_input(input_config):
- input_path = input_config['input_path']
- image = cv2.imread(input_path, cv2.IMREAD_COLOR)
- h, w, _ = image.shape
- img_scale = (400, 400)
- max_long_edge = max(img_scale)
- max_short_edge = min(img_scale)
- scale_factor = min(max_long_edge / max(h, w), max_short_edge / min(h, w))
- scale_w = int(w * float(scale_factor) + 0.5)
- scale_h = int(h * float(scale_factor) + 0.5)
- image = cv2.resize(image, (scale_w, scale_h))
-
- image = np.asarray(image).astype(np.float32)
- mean = np.array([123.675, 116.28, 103.53], dtype=np.float32)
- std = np.array([58.395, 57.12, 57.375], dtype=np.float32)
- image = imnormalize(image, mean, std)
-
- divisor = 32
- pad_h = int(np.ceil(image.shape[0] / divisor)) * divisor
- pad_w = int(np.ceil(image.shape[1] / divisor)) * divisor
- image = impad(image, shape=(pad_h, pad_w), pad_val=0)
- image = np.asarray(image).astype(np.float32)
-
- image = np.transpose(image, [2, 0, 1])
- one_img = torch.from_numpy(image).unsqueeze(0).float().requires_grad_(
- True)
- (_, C, H, W) = input_config['input_shape']
- one_meta = {
- 'img_shape': (H, W, C),
- 'ori_shape': (H, W, C),
- 'pad_shape': (H, W, C),
- 'filename': '<demo>.png',
- 'scale_factor': np.ones(4, dtype=np.float32),
- 'flip': False,
- 'show_img': image,
- 'flip_direction': None
- }
-
- return one_img, one_meta
- print(f"onnxruntime device: {ort.get_device()}") # output: GPU
- print(f'ort avail providers: {ort.get_available_providers()}') # output: ['CUDAExecutionProvider', 'CPUExecutionProvider']
- output_file = "checkpoints/AD_dsxw_atts_20220318.onnx"
- classes_name = ['yiwei','loujian','celi','libei','fantie','lianxi','duojian','shunjian','shaoxi','jiahan','yiwu']
- onnx_model = onnx.load(output_file)
- onnx.checker.check_model(onnx_model)
- onnx_model = ONNXRuntimeDetector(output_file, classes_name, 0)
- input_shape = (1, 3, 400, 400)
- input_img = "/home/shanwei-luo/userdata/datasets/dsxw_dataset_v4/dsxw_test/images/21000204.jpg"
- normalize_cfg = {'mean': (123.675, 116.28, 103.53),'std': (58.395, 57.12, 57.375)}
- input_config = {
- 'input_shape': input_shape,
- 'input_path': input_img,
- 'normalize_cfg': normalize_cfg
- }
-
- input_config['input_shape'] = (1, 3, 416, 416)
-
- # prepare input once again
- one_img, one_meta = preprocess_example_input(input_config)
- print(one_img)
- img_list, img_meta_list = [one_img], [[one_meta]]
- img_list = [_.cuda().contiguous() for _ in img_list]
- onnx_results = onnx_model(
- img_list, img_metas=img_meta_list, return_loss=False)[0]
- print(onnx_results)
-
- config_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test61/AD_dsxw_test61.py'
- checkpoint_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test61/epoch_36.pth'
- model_1 = init_detector(config_file_1, checkpoint_file_1, device='cuda:1')
- results_1_tmp = inference_detector(model_1, [input_img])
- print(results_1_tmp)
-
- for i in range(len(onnx_results)):
- print(onnx_results[i].shape)
- for i in range(len(results_1_tmp[0])):
- print(len(results_1_tmp[0][i]))
|