def main():
args = parse_args()
cfg = mmcv.Config.fromfile(args.config)
cfg.data.test.test_mode = True
trt_model = init_detector(args.trt_model_path)
# create dataset and dataloader
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(
dataset,
samples_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=False,
shuffle=False)
model = ModelWarper(trt_model, num_classes=len(dataset.CLASSES))
outputs = single_gpu_test(model, data_loader)
if args.out:
print('\nwriting results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
# outputs = mmcv.load(args.out)
import numpy as np
print('eval bbox:')
eval_result=dataset.evaluate(outputs, metric='bbox', classwise=False)
print(eval_result)
def model_test(test_folder,
cfg_path,
checkpoint,
save_folder,
opt_shape_param=None,
max_workspace_size=1 << 25,
device="cuda:0",
score_thr=0.3,
fp16=True,
enable_mask=False):
if not osp.exists(save_folder):
os.mkdir(save_folder)
trt_model_path = osp.join(save_folder, 'trt_model.pth')
logger.info("creating {} trt model.".format(cfg_path))
trt_model = mmdet2trt(cfg_path,
checkpoint,
opt_shape_param=opt_shape_param,
max_workspace_size=int(max_workspace_size),
fp16_mode=fp16,
device=device,
enable_mask=enable_mask)
logger.info("finish, save trt_model in {}".format(trt_model_path))
torch.save(trt_model.state_dict(), trt_model_path)
trt_model = init_detector(trt_model_path)
file_list = os.listdir(test_folder)
for file_name in tqdm.tqdm(file_list):
if not file_name.lower().endswith('.jpg') or file_name.lower(
).endswith('.png'):
continue
image_path = osp.join(test_folder, file_name)
result = inference_detector(trt_model, image_path, cfg_path, device)
num_detections = result[0].item()
trt_bbox = result[1][0]
trt_score = result[2][0]
trt_cls = result[3][0]
image = cv2.imread(image_path)
for i in range(num_detections):
scores = trt_score[i].item()
classes = int(trt_cls[i].item())
if scores < score_thr:
continue
bbox = tuple(trt_bbox[i])
bbox = tuple(int(v) for v in bbox)
color = ((classes >> 2 & 1) * 128 + (classes >> 5 & 1) * 128,
(classes >> 1 & 1) * 128 + (classes >> 4 & 1) * 128,
(classes >> 0 & 1) * 128 + (classes >> 3 & 1) * 128)
cv2.rectangle(image, bbox[:2], bbox[2:], color, thickness=5)
cv2.imwrite(osp.join(save_folder, file_name), image)
def main():
parser = ArgumentParser()
parser.add_argument('img', help='Image file')
parser.add_argument('config', help='mmdet Config file')
parser.add_argument('checkpoint', help='mmdet Checkpoint file')
parser.add_argument('save_path', help='tensorrt model save path')
parser.add_argument('--device',
default='cuda:0',
help='Device used for inference')
parser.add_argument('--score-thr',
type=float,
default=0.3,
help='bbox score threshold')
parser.add_argument("--fp16",
type=bool,
default=True,
help="enable fp16 inference")
args = parser.parse_args()
cfg_path = args.config
trt_model = mmdet2trt(cfg_path,
args.checkpoint,
fp16_mode=args.fp16,
device=args.device)
torch.save(trt_model.state_dict(), args.save_path)
trt_model = init_detector(args.save_path)
image_path = args.img
result = inference_detector(trt_model, image_path, cfg_path, args.device)
num_detections = result[0].item()
trt_bbox = result[1][0]
trt_score = result[2][0]
trt_cls = result[3][0]
image = cv2.imread(image_path)
input_image_shape = image.shape
for i in range(num_detections):
scores = trt_score[i].item()
classes = int(trt_cls[i].item())
if scores < args.score_thr:
continue
bbox = tuple(trt_bbox[i])
bbox = tuple(int(v) for v in bbox)
color = ((classes >> 2 & 1) * 128 + (classes >> 5 & 1) * 128,
(classes >> 1 & 1) * 128 + (classes >> 4 & 1) * 128,
(classes >> 0 & 1) * 128 + (classes >> 3 & 1) * 128)
cv2.rectangle(image, bbox[:2], bbox[2:], color, thickness=5)
if input_image_shape[0] > 1280 or input_image_shape[1] > 720:
scales = min(720 / image.shape[0], 1280 / image.shape[1])
image = cv2.resize(image, (0, 0), fx=scales, fy=scales)
cv2.imshow('image', image)
c = cv2.waitKey()
def main():
parser = ArgumentParser()
parser.add_argument('test_folder', help='folder contain test images')
parser.add_argument('config', help='mmdet Config file')
parser.add_argument('checkpoint', help='mmdet Checkpoint file')
parser.add_argument(
'save_folder',
help='tensorrt model and test images results save folder')
parser.add_argument('--trt_model_path',
default='',
help='save and inference model. '
'default [save_folder]/trt_model.pth')
parser.add_argument(
'--opt_shape_param',
default='[ [ [1,3,800,800], [1,3,800,1344], [1,3,1344,1344] ] ]',
help='min/opt/max shape of input')
parser.add_argument('--max_workspace_size',
default=1 << 30,
help='max workspace size')
parser.add_argument('--device',
default='cuda:0',
help='Device used for inference')
parser.add_argument('--score-thr',
type=float,
default=0.3,
help='bbox score threshold')
parser.add_argument('--fp16',
action='store_true',
help='enable fp16 inference')
parser.add_argument('--enable_mask',
action='store_true',
help='enable mask output')
parser.add_argument('--test-mode',
default='all',
help='what to do in the test',
choices=['convert', 'inference', 'all'])
args = parser.parse_args()
trt_model_path = args.trt_model_path
if len(trt_model_path) == 0:
trt_model_path = osp.join(args.save_folder, 'test_model.pth')
if not osp.exists(args.save_folder):
os.mkdir(args.save_folder)
test_mode = TEST_MODE_DICT[args.test_mode]
if test_mode & TEST_MODE_DICT['convert'] > 0:
convert_test(args.config,
args.checkpoint,
trt_model_path,
opt_shape_param=eval(args.opt_shape_param),
max_workspace_size=args.max_workspace_size,
device=args.device,
fp16=args.fp16)
trt_model = init_detector(trt_model_path)
else:
trt_model = init_detector(trt_model_path)
if test_mode & TEST_MODE_DICT['inference'] > 0:
inference_test(trt_model,
args.config,
args.device,
args.test_folder,
args.save_folder,
score_thr=args.score_thr)
def main():
parser = ArgumentParser()
parser.add_argument('img', help='Image file')
parser.add_argument('config', help='mmdet Config file')
parser.add_argument('checkpoint', help='mmdet Checkpoint file')
parser.add_argument('save_path', help='tensorrt model save path')
parser.add_argument('--device',
default='cuda:0',
help='Device used for inference')
parser.add_argument('--score-thr',
type=float,
default=0.3,
help='bbox score threshold')
parser.add_argument('--fp16',
action='store_true',
help="enable fp16 inference")
args = parser.parse_args()
cfg_path = args.config
opt_shape_param = None
if opt_shape_param is None:
#img_scale = cfg_path.test_pipeline[1]['img_scale']
#min_scale = min(img_scale)
#max_scale = max(img_scale)
#print(min_scale)
opt_shape_param=[
[
[1,3,224,224], # min tensor shape
[1,3,800,1312], # shape used to do int8 calib
[1,3,1344,1344], # max tensor shape
]
]
trt_model = mmdet2trt(cfg_path,
args.checkpoint,
opt_shape_param=opt_shape_param,
fp16_mode=args.fp16,
device=args.device)
torch.save(trt_model.state_dict(), args.save_path)
trt_model = init_detector(args.save_path)
image_path = args.img
result = inference_detector(trt_model, image_path, cfg_path, args.device)
num_detections = result[0].item()
trt_bbox = result[1][0]
trt_score = result[2][0]
trt_cls = result[3][0]
image = cv2.imread(image_path)
input_image_shape = image.shape
for i in range(num_detections):
scores = trt_score[i].item()
classes = int(trt_cls[i].item())
if scores < args.score_thr:
continue
bbox = tuple(trt_bbox[i])
bbox = tuple(int(v) for v in bbox)
color = ((classes >> 2 & 1) * 128 + (classes >> 5 & 1) * 128,
(classes >> 1 & 1) * 128 + (classes >> 4 & 1) * 128,
(classes >> 0 & 1) * 128 + (classes >> 3 & 1) * 128)
cv2.rectangle(image, bbox[:2], bbox[2:], color, thickness=5)
if input_image_shape[0] > 1280 or input_image_shape[1] > 720:
scales = min(720 / image.shape[0], 1280 / image.shape[1])
image = cv2.resize(image, (0, 0), fx=scales, fy=scales)
cv2.imwrite('image', image)
#c = cv2.waitKey()
if __name__ == '__main__':
main()