from mmdet.apis import init_detector, inference_detector

"""
config_file = 'configs/swin/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_3x_coco.py'
# 从 model zoo 下载 checkpoint 并放在 `checkpoints/` 文件下 网址为:
# http://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth
checkpoint_file = 'mask_rcnn_swin_tiny_patch4_window7.pth'
device = 'cuda:0'
# 初始化检测器
model = init_detector(config_file, checkpoint_file, device=device)
# 推理演示图像
inference_detector(model, 'demo/demo.jpg')
"""
from argparse import ArgumentParser
from mmdet.apis import inference_detector, init_detector, show_result_pyplot
import cv2 as cv

# config_file = 'configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py'
config_file = 'configs/swin/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_3x_coco.py'
# download the checkpoint from model zoo and put it in `checkpoints/`
# url: https://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth
# checkpoint_file = 'faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth'
checkpoint_file = 'checkpoints/epoch_200.pth'
device = 'cuda:0'
# init a detector
model = init_detector(config_file, checkpoint_file, device=device)
# inference the demo image
img = 'demo/31.jpg'
result = inference_detector(model, img)
# model.show_result(img, result)
# cv.waitkey(1000)
# model.show_result(img, result, out_file='demo/demo_result.jpg') # 保存推理图像
show_result_pyplot(model, img, result)
model.show_result(img, result, out_file='demo/demo_31_result2.jpg')  # 保存推理图像