Spaces:
Sleeping
Sleeping
| # Copyright (c) OpenMMLab. All rights reserved. | |
| """Deploy demo for mmdeploy. | |
| This script help user to run mmdeploy demo after convert the | |
| checkpoint to backends. | |
| Usage: | |
| python deploy_demo.py img \ | |
| config \ | |
| checkpoint \ | |
| [--deploy-cfg DEPLOY_CFG] \ | |
| [--device DEVICE] \ | |
| [--out-dir OUT_DIR] \ | |
| [--show] \ | |
| [--score-thr SCORE_THR] | |
| Example: | |
| python deploy_demo.py \ | |
| ${MMYOLO_PATH}/data/cat/images \ | |
| ./yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py \ | |
| ./end2end.engine \ | |
| --deploy-cfg ./detection_tensorrt-fp16_dynamic-192x192-960x960.py \ | |
| --out-dir ${MMYOLO_PATH}/work_dirs/deploy_predict_out \ | |
| --device cuda:0 \ | |
| --score-thr 0.5 | |
| """ | |
| import argparse | |
| import os | |
| import torch | |
| from mmengine import ProgressBar | |
| from mmyolo.utils.misc import get_file_list | |
| try: | |
| from mmdeploy.apis.utils import build_task_processor | |
| from mmdeploy.utils import get_input_shape, load_config | |
| except ImportError: | |
| raise ImportError( | |
| 'mmdeploy is not installed, please see ' | |
| 'https://mmdeploy.readthedocs.io/en/1.x/01-how-to-build/build_from_source.html' # noqa | |
| ) | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description='For mmdeploy predict') | |
| parser.add_argument( | |
| 'img', help='Image path, include image file, dir and URL.') | |
| parser.add_argument('config', help='model config root') | |
| parser.add_argument('checkpoint', help='checkpoint backend model path') | |
| parser.add_argument('--deploy-cfg', help='deploy config path') | |
| parser.add_argument( | |
| '--device', default='cuda:0', help='device used for conversion') | |
| parser.add_argument( | |
| '--out-dir', default='./output', help='Path to output file') | |
| parser.add_argument( | |
| '--show', action='store_true', help='Show the detection results') | |
| parser.add_argument( | |
| '--score-thr', type=float, default=0.3, help='Bbox score threshold') | |
| args = parser.parse_args() | |
| return args | |
| # TODO Still need to refactor to not building dataset. | |
| def main(): | |
| args = parse_args() | |
| if not os.path.exists(args.out_dir) and not args.show: | |
| os.mkdir(args.out_dir) | |
| # read deploy_cfg and config | |
| deploy_cfg, model_cfg = load_config(args.deploy_cfg, args.config) | |
| # build task and backend model | |
| task_processor = build_task_processor(model_cfg, deploy_cfg, args.device) | |
| model = task_processor.build_backend_model([args.checkpoint]) | |
| # get model input shape | |
| input_shape = get_input_shape(deploy_cfg) | |
| # get file list | |
| files, source_type = get_file_list(args.img) | |
| # start detector inference | |
| progress_bar = ProgressBar(len(files)) | |
| for file in files: | |
| # process input image | |
| model_inputs, _ = task_processor.create_input(file, input_shape) | |
| # do model inference | |
| with torch.no_grad(): | |
| result = model.test_step(model_inputs) | |
| if source_type['is_dir']: | |
| filename = os.path.relpath(file, args.img).replace('/', '_') | |
| else: | |
| filename = os.path.basename(file) | |
| out_file = None if args.show else os.path.join(args.out_dir, filename) | |
| # filter score | |
| result = result[0] | |
| result.pred_instances = result.pred_instances[ | |
| result.pred_instances.scores > args.score_thr] | |
| # visualize results | |
| task_processor.visualize( | |
| image=file, | |
| model=model, | |
| result=result, | |
| show_result=args.show, | |
| window_name=os.path.basename(filename), | |
| output_file=out_file) | |
| progress_bar.update() | |
| print('All done!') | |
| if __name__ == '__main__': | |
| main() | |