| | from argparse import ArgumentParser |
| |
|
| | import numpy as np |
| | import requests |
| |
|
| | from mmdet.apis import inference_detector, init_detector, show_result_pyplot |
| | from mmdet.core import bbox2result |
| |
|
| |
|
| | def parse_args(): |
| | parser = ArgumentParser() |
| | parser.add_argument('img', help='Image file') |
| | parser.add_argument('config', help='Config file') |
| | parser.add_argument('checkpoint', help='Checkpoint file') |
| | parser.add_argument('model_name', help='The model name in the server') |
| | parser.add_argument( |
| | '--inference-addr', |
| | default='127.0.0.1:8080', |
| | help='Address and port of the inference server') |
| | parser.add_argument( |
| | '--device', default='cuda:0', help='Device used for inference') |
| | parser.add_argument( |
| | '--score-thr', type=float, default=0.5, help='bbox score threshold') |
| | args = parser.parse_args() |
| | return args |
| |
|
| |
|
| | def parse_result(input, model_class): |
| | bbox = [] |
| | label = [] |
| | score = [] |
| | for anchor in input: |
| | bbox.append(anchor['bbox']) |
| | label.append(model_class.index(anchor['class_name'])) |
| | score.append([anchor['score']]) |
| | bboxes = np.append(bbox, score, axis=1) |
| | labels = np.array(label) |
| | result = bbox2result(bboxes, labels, len(model_class)) |
| | return result |
| |
|
| |
|
| | def main(args): |
| | |
| | model = init_detector(args.config, args.checkpoint, device=args.device) |
| | |
| | model_result = inference_detector(model, args.img) |
| | for i, anchor_set in enumerate(model_result): |
| | anchor_set = anchor_set[anchor_set[:, 4] >= 0.5] |
| | model_result[i] = anchor_set |
| | |
| | show_result_pyplot( |
| | model, |
| | args.img, |
| | model_result, |
| | score_thr=args.score_thr, |
| | title='pytorch_result') |
| | url = 'http://' + args.inference_addr + '/predictions/' + args.model_name |
| | with open(args.img, 'rb') as image: |
| | response = requests.post(url, image) |
| | server_result = parse_result(response.json(), model.CLASSES) |
| | show_result_pyplot( |
| | model, |
| | args.img, |
| | server_result, |
| | score_thr=args.score_thr, |
| | title='server_result') |
| |
|
| | for i in range(len(model.CLASSES)): |
| | assert np.allclose(model_result[i], server_result[i]) |
| |
|
| |
|
| | if __name__ == '__main__': |
| | args = parse_args() |
| | main(args) |
| |
|
