def make_trackvideo(configpath, weightpath, imagefolder, videofolder):
os.makedirs(videofolder, exist_ok=True)
imagenames = sorted(os.listdir(imagefolder))
# find video names
videonames = set()
for img_nm in imagenames:
video = img_nm.rsplit('_', 1)[0]
videonames.add(video)
print(f"video files: {videonames}\n")
detector = init_detector(configpath, weightpath)
# make video files
for vdo_id, vdo_nm in enumerate(videonames):
fourcc = cv2.VideoWriter_fourcc('M','J','P','G')
fps = 12
shape = (800, 800)
vdo_pth = os.path.join(videofolder, f"{vdo_nm}.avi")
result_video = cv2.VideoWriter(vdo_pth, fourcc, fps, shape)
image_paths = glob.glob(os.path.join(imagefolder, f"{vdo_nm}*"))
image_paths = sorted(image_paths, key=lambda x: int(os.path.splitext(x)[0].rsplit('_', 1)[1]))
tracker = Sort()
for img_id, img_pth in enumerate(image_paths):
image = cv2.imread(img_pth)
result, inf_time = inference_detector(detector, image)
detected_boxes = result[0]
tracked_boxes = tracker.update(detected_boxes)
# draw the result on image
bboxes, track_ids = tracked_boxes[:, :-1], tracked_boxes[:, -1]
for box, trk_id in zip(bboxes, track_ids):
x1, y1, x2, y2 = box.astype("int")
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 2)
cx, cy = (x1+x2)//2, (y1+y2)//2
text = f"{trk_id.astype(int)}"
cv2.putText(image, text, (cx-10, cy-10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
cv2.circle(image, (cx, cy), 4, (0, 0, 255), -1)
image = cv2.resize(image, shape)
result_video.write(image)
print(f"video {vdo_id} frame {img_id} drawing done.")
result_video.release()
print(f"\n{vdo_nm}.avi created!!\n")
return None
def main():
args = parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = '6'
model = init_detector(args.config,
args.model,
device=torch.device('cuda', args.gpu))
#root_dir = "/tcdata/guangdong1_round2_testA_20190924"
#root_dir = "/data/guangdong1_round2_testA_20190924"
root_dir = "data/defect"
image_dirs = os.listdir(root_dir)
count = len(image_dirs)
results = []
small_gt_count = 0
for image_index, image_dir in enumerate(image_dirs):
if not os.path.isdir(os.path.join(root_dir, image_dir)):
continue
image_files = os.listdir(os.path.join(root_dir, image_dir))
image_path = os.path.join(root_dir, image_dir, image_dir + '.jpg')
img = cv2.imread(image_path)
result = inference_detector(model, img)
for index, bboxes in enumerate(result):
for bbox in bboxes:
score = bbox[4].item()
cls_type = get_cls_type(index)
new_bbox = [
round(bbox[0].item(), 2),
round(bbox[1].item(), 2),
round(bbox[2].item(), 2),
round(bbox[3].item(), 2)
]
if new_bbox[2] - new_bbox[0] < 16 or new_bbox[3] - new_bbox[
1] < 16:
small_gt_count += 1
print(
"width = {}, height = {}, small_gt_count = {}, img = {}"
.format(new_bbox[2] - new_bbox[0],
new_bbox[3] - new_bbox[1], small_gt_count,
image_dir))
name = image_dir + '.jpg'
result = {
'name': name,
'category': cls_type,
'bbox': new_bbox,
'score': score
}
results.append(result)
#print("\r"+"{}/{}".format(image_index, count), end="", flush=True)
with open('../result.json', 'w') as fp:
json.dump(results, fp, indent=4, separators=(',', ': '))
def __init__(self,
config_file,
checkpoint_file,
image_dir=None,
labels_file=None,
score_threshold=0.3,
device='cpu',
**kwargs):
"""
Load MMDetection model from config and checkpoint into memory.
(Check https://mmdetection.readthedocs.io/en/v1.2.0/GETTING_STARTED.html#high-level-apis-for-testing-images)
Optionally set mappings from COCO classes to target labels
:param config_file: Absolute path to MMDetection config file (e.g. /home/user/mmdetection/configs/faster_rcnn/faster_rcnn_r50_fpn_1x.py)
:param checkpoint_file: Absolute path MMDetection checkpoint file (e.g. /home/user/mmdetection/checkpoints/faster_rcnn_r50_fpn_1x_20181010-3d1b3351.pth)
:param image_dir: Directory where images are stored (should be used only in case you use direct file upload into Label Studio instead of URLs)
:param labels_file: file with mappings from COCO labels to custom labels {"airplane": "Boeing"}
:param score_threshold: score threshold to wipe out noisy results
:param device: device (cpu, cuda:0, cuda:1, ...)
:param kwargs:
"""
super(MMDetection, self).__init__(**kwargs)
self.config_file = config_file
self.checkpoint_file = checkpoint_file
self.labels_file = labels_file
# default Label Studio image upload folder
upload_dir = os.path.join(get_data_dir(), 'media', 'upload')
self.image_dir = image_dir or upload_dir
logger.debug(
f'{self.__class__.__name__} reads images from {self.image_dir}')
if self.labels_file and os.path.exists(self.labels_file):
self.label_map = json_load(self.labels_file)
else:
self.label_map = {}
self.from_name, self.to_name, self.value, self.labels_in_config = get_single_tag_keys(
self.parsed_label_config, 'RectangleLabels', 'Image')
schema = list(self.parsed_label_config.values())[0]
self.labels_in_config = set(self.labels_in_config)
# Collect label maps from `predicted_values="airplane,car"` attribute in <Label> tag
self.labels_attrs = schema.get('labels_attrs')
if self.labels_attrs:
for label_name, label_attrs in self.labels_attrs.items():
for predicted_value in label_attrs.get('predicted_values',
'').split(','):
self.label_map[predicted_value] = label_name
print('Load new model from: ', config_file, checkpoint_file)
self.model = init_detector(config_file, checkpoint_file, device=device)
self.score_thresh = score_threshold
def test_and_draw_from_single_file(sample_file, ext_name, bgr_file, out_file, config_file, checkpoint_file, score_threhold):
model = init_detector(config_file, checkpoint_file, device='cuda:0')
sample = None
if ext_name == 'bgr':
sample = mmcv.imread(sample_file)
elif ext_name == 'tiff':
sample = load_pol_sub_image(sample_file)
else:
sample = LoadPolNPZImageFromFile(sample_file)
print(sample)
result = inference_detector(model, sample)
img = mmcv.imread(bgr_file)
show_result(img, result, model.CLASSES, out_file=out_file, score_thr=score_threhold)
async def async_main(args):
# build the model from a config file and a checkpoint file
model = init_detector(args.config, args.checkpoint, device=args.device)
# test a single image
tasks = asyncio.create_task(async_inference_detector(model, args.img))
result = await asyncio.gather(tasks)
# show the results
show_result_pyplot(model,
args.img,
result[0],
palette=args.palette,
score_thr=args.score_thr,
out_file=args.out_file)
def __init__(self, mode):
self.config_file = './configs/mask_rcnn/mask_rcnn_r50_fpn_1x_coco.py'
self.checkpoint_file = '../common_data/mask_rcnn_r50_fpn_1x_coco_20200205-d4b0c5d6.pth'
self.Mask_RCNN = init_detector(self.config_file,
self.checkpoint_file,
device='cpu').cuda()
self.mean = torch.Tensor(
[123.675, 116.28,
103.53]).unsqueeze(0).unsqueeze(-1).unsqueeze(-1).cuda()
self.std = torch.Tensor(
[58.395, 57.12,
57.375]).unsqueeze(0).unsqueeze(-1).unsqueeze(-1).cuda()
def __init__(self,
model_config,
checkpoint=None,
streamqueue_size=3,
device='cuda:0'):
self.streamqueue_size = streamqueue_size
self.device = device
# build the model and load checkpoint
self.model = init_detector(model_config,
checkpoint=None,
device=self.device)
self.streamqueue = None
def main(args):
# build the model from a config file and a checkpoint file
model = init_detector(args.config, args.checkpoint, device=args.device)
# test a single image
result = inference_detector(model, args.img)
# show the results
show_result_pyplot(model, args.img, result, score_thr=args.score_thr)
# добавлено сохранение в файл
model.show_result(args.img,
result,
score_thr=args.score_thr,
font_size=5,
out_file='result_.jpg')
def main(args):
# build the model from a config file and a checkpoint file
model = init_detector(args.config, args.checkpoint, device=args.device)
# test a single image
result = inference_detector(model, args.img)
# show the results
show_result_pyplot(
model,
args.img,
result,
palette=args.palette,
score_thr=args.score_thr,
out_file=args.out_file)
def main():
args = parse_args()
model = init_detector(
args.config, args.checkpoint, device=torch.device('cuda', args.device))
camera = cv2.VideoCapture(args.camera_id)
while True:
ret_val, img = camera.read()
result = inference_detector(model, img)
show_result(
img, result, model.CLASSES, score_thr=args.score_thr, wait_time=1)
def __init__(self, model_path_1,model_path_2, config_path, json_out_path, pic_path):
self.init_start = time.time()
self.model1 = None
self.model2 = init_detector(config_path, model_path_2, device='cuda:0')
self.pics = glob.glob(os.path.join(pic_path, 'img_*'))
self.bottle_pics = np.unique([name.split('_')[-2] for name in glob.glob(os.path.join(pic_path, 'imgs_*'))])
self.json_out_path = json_out_path
self.pic_path = pic_path
self.images = []
self.annotations = []
self.num_img_id = 1
self.num_ann_id = 1
self.init_end = time.time()
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--config_file',
type=str,
default=
'/home/songbai.xb/detection/mmdetection/myprojects/extinguisher/configs/reppoints_minmax_r50_fpn_1x.py',
help='config_file')
parser.add_argument(
'--checkpoint_file',
type=str,
default=
'/home/songbai.xb/detection/mmdetection/myprojects/extinguisher/work_dirs/all/reppoints_minmax_r50_fpn_mykeep_ratio_RandAugment/latest.pth',
help='model weight')
parser.add_argument(
'--export_file',
type=str,
default=
'/home/songbai.xb/detection/mmdetection_merge/mmdetection/myprojects/extinguisher/reppoints_minmax_r50_fpn_1x.pt',
help='the file to save export model')
parser.add_argument('--work_size',
type=int,
nargs=2,
metavar=('height', 'width'),
help='model input size')
parser.add_argument('--gpuid',
type=str,
default='0',
help='visible gpu ids')
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpuid
config_file = args.config_file
checkpoint_file = args.checkpoint_file
# detector init
detector = init_detector(config_file, checkpoint_file, device='cuda:0')
# export detector
script_module = export_model(detector, args.work_size, args.export_file)
# check export detector
check_jit_model(detector, args.work_size, args.export_file)
def main(args):
assert args.opset in available_opsets
assert args.opset > 9
torch.set_default_tensor_type(torch.FloatTensor)
model = init_detector(args.config, args.checkpoint, device='cpu')
model.eval()
if torch.cuda.is_available():
model.cuda()
device = next(model.parameters()).device
cfg = model.cfg
fake_data = get_fake_input(cfg, device=device)
# BEGIN nncf part
if cfg.get('nncf_config'):
check_nncf_is_enabled()
if not is_checkpoint_nncf(args.checkpoint):
raise RuntimeError('Trying to make export with NNCF compression '
'a model snapshot that was NOT trained with NNCF')
cfg.load_from = args.checkpoint
cfg.resume_from = None
compression_ctrl, model = wrap_nncf_model(model, cfg, None, get_fake_input)
# TODO: apply the following string for NNCF 1.5.*
#compression_ctrl.prepare_for_export()
# END nncf part
if args.target == 'openvino' and not args.alt_ssd_export:
if hasattr(model, 'roi_head'):
stub_roi_feature_extractor(model.roi_head, 'bbox_roi_extractor')
stub_roi_feature_extractor(model.roi_head, 'mask_roi_extractor')
mmcv.mkdir_or_exist(osp.abspath(args.output_dir))
onnx_model_path = osp.join(args.output_dir,
osp.splitext(osp.basename(args.config))[0] + '.onnx')
with torch.no_grad():
export_to_onnx(model, fake_data, export_name=onnx_model_path, opset=args.opset,
alt_ssd_export=getattr(args, 'alt_ssd_export', False),
verbose=True)
add_node_names(onnx_model_path)
print(f'ONNX model has been saved to "{onnx_model_path}"')
optimize_onnx_graph(onnx_model_path)
if args.target == 'openvino':
input_shape = list(fake_data['img'][0].shape)
if args.input_shape:
input_shape = [1, 3, *args.input_shape]
export_to_openvino(cfg, onnx_model_path, args.output_dir, input_shape, args.input_format)
else:
pass
def TestInfer(score_thr=0.8):
vin_test_file = '/data/pycode/CXRAD/dataset/VinCXR_test.txt'
vin_test_image = '/data/fjsdata/Vin-CXR/test_jpg/'
vin_test_data = '/data/comcode/mmdetection/vincxr/test/'
# Specify the path to model config and checkpoint file
config_file = 'vincxr/code/maskrcnn.py'
checkpoint_file = 'vincxr/workdir/latest.pth'
# build the model from a config file and a checkpoint file
model = init_detector(config_file, checkpoint_file, device='cuda:6')
# test images and show the results
images = pd.read_csv(vin_test_file, sep=',', header=None).values
sub_res = []
for image in images:
img = vin_test_image + image[0]+'.jpeg'
result = inference_detector(model, img)
#extract result
if isinstance(result, tuple):
bbox_result, segm_result = result
if isinstance(segm_result, tuple):
segm_result = segm_result[0] # ms rcnn
else:
bbox_result, segm_result = result, None
bboxes = np.vstack(bbox_result)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
#prediction
assert bboxes.shape[1] == 5
scores = bboxes[:, -1]
sub_tmp = {'image_id': image[0], 'PredictionString': '14 1.0 0 0 1 1'}
if len(scores)>0:
inds = scores > score_thr
bboxes = bboxes[inds, :]
labels = labels[inds]
scores = scores[inds]
if len(scores)>0:
sub_tmp['PredictionString'] = format_prediction_string(labels, bboxes, scores)
#sub = {'image_id': image[0],'PredictionString': format_prediction_string(labels, bboxes, scores)}
sub_res.append(sub_tmp)
sys.stdout.write('\r process: = {}'.format(len(sub_res)))
sys.stdout.flush()
#Save submission file
test_df = pd.DataFrame(sub_res, columns=['image_id', 'PredictionString'])
print("\r set shape: {}".format(test_df.shape))
print("\r set Columns: {}".format(test_df.columns))
test_df.to_csv(vin_test_data+'submission.csv', index=False)
def __init__(self,
config_file,
checkpoint_file,
labels_file=None,
score_threshold=0.3,
device="cpu",
**kwargs):
"""
Load MMDetection model from config and checkpoint into memory.
(Check https://mmdetection.readthedocs.io/en/v1.2.0/GETTING_STARTED.html#high-level-apis-for-testing-images)
Optionally set mappings from COCO classes to target labels
:param config_file: Absolute path to MMDetection config file (e.g. /home/user/mmdetection/configs/faster_rcnn/faster_rcnn_r50_fpn_1x.py)
:param checkpoint_file: Absolute path MMDetection checkpoint file (e.g. /home/user/mmdetection/checkpoints/faster_rcnn_r50_fpn_1x_20181010-3d1b3351.pth)
:param labels_file: file with mappings from COCO labels to custom labels {"airplane": "Boeing"}
:param score_threshold: score threshold to wipe out noisy results
:param device: device (cpu, cuda:0, cuda:1, ...)
:param kwargs:
"""
super(MMDetection, self).__init__(**kwargs)
self.config_file = config_file
self.checkpoint_file = checkpoint_file
self.labels_file = labels_file
if self.labels_file and os.path.exists(self.labels_file):
self.label_map = json_load(self.labels_file)
else:
self.label_map = {}
(
self.from_name,
self.to_name,
self.value,
self.labels_in_config,
) = get_single_tag_keys(self.parsed_label_config, "RectangleLabels",
"Image")
schema = list(self.parsed_label_config.values())[0]
self.labels_in_config = set(self.labels_in_config)
# Collect label maps from `predicted_values="airplane,car"` attribute in <Label> tag
self.labels_attrs = schema.get("labels_attrs")
if self.labels_attrs:
for label_name, label_attrs in self.labels_attrs.items():
for predicted_value in label_attrs.get("predicted_values",
"").split(","):
self.label_map[predicted_value] = label_name
print("Load new model from: ", config_file, checkpoint_file)
self.model = init_detector(config_file, checkpoint_file, device=device)
self.score_thresh = score_threshold
def detection_inference(args, frame_paths):
model = init_detector(args.det_config, args.det_checkpoint, args.device)
assert model.CLASSES[0] == 'person', ('We require you to use a detector '
'trained on COCO')
results = []
print('Performing Human Detection for each frame')
prog_bar = mmcv.ProgressBar(len(frame_paths))
for frame_path in frame_paths:
result = inference_detector(model, frame_path)
# We only keep human detections with score larger than det_score_thr
result = result[0][result[0][:, 4] >= args.det_score_thr]
results.append(result)
prog_bar.update()
return results
def model_init(cfg_files=None, ckpt_files=None):
if cfg_files == None:
cfg_files = CFG_FILES
if ckpt_files == None:
ckpt_files = CKPT_FILES
if not isinstance(cfg_files, list):
cfg_files = [cfg_files]
if not isinstance(ckpt_files, list):
ckpt_files = [ckpt_files]
models = []
for idx, (cfg, ckp) in enumerate(zip(cfg_files, ckpt_files)):
model = init_detector(cfg, ckp, device='cuda:0')
models.append(model)
return models
def main():
args = parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
model = init_detector(args.config,
args.model,
device=torch.device('cuda', args.gpu))
root_dir = "data/normal"
#root_dir = "/tcdata/guangdong1_round2_testA_20190924"
#root_dir = "/data/guangdong1_round2_testA_20190924"
image_dirs = os.listdir(root_dir)
count = len(image_dirs)
results = []
for image_index, image_dir in enumerate(image_dirs):
if not os.path.isdir(os.path.join(root_dir, image_dir)):
continue
image_files = os.listdir(os.path.join(root_dir, image_dir))
image_path = os.path.join(root_dir, image_dir, image_dir + '.jpg')
img = cv2.imread(image_path)
style = image_dir.split('_')[0]
template_path = os.path.join(root_dir, image_dir,
"template_{}.jpg".format(style))
template_img = cv2.imread(template_path)
result = inference_detector(model, img, template_img)
for index, bboxes in enumerate(result):
if index >= 15:
continue
for bbox in bboxes:
score = bbox[4].item()
cls_type = get_cls_type(index)
new_bbox = [
round(bbox[0].item(), 2),
round(bbox[1].item(), 2),
round(bbox[2].item(), 2),
round(bbox[3].item(), 2)
]
#print("class type = {}, bbox = {}, score = {}".format(cls_type, bbox, score))
name = image_dir + '.jpg'
result = {
'name': name,
'category': cls_type,
'bbox': new_bbox,
'score': score
}
results.append(result)
print("\r" + "{}/{}".format(image_index, count), end="", flush=True)
with open('../result.json', 'w') as fp:
json.dump(results, fp, indent=4, separators=(',', ': '))
def main():
parser = ArgumentParser()
parser.add_argument('--imgs', help='Image file')
parser.add_argument('--output_dir')
parser.add_argument('--config', help='Config file')
parser.add_argument('--checkpoint', help='Checkpoint file')
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')
args = parser.parse_args()
# build the model from a config file and a checkpoint file
model = init_detector(args.config, args.checkpoint, device=args.device)
os.makedirs(os.path.join(args.output_dir, 'vis'), exist_ok=True)
os.makedirs(os.path.join(args.output_dir, 'crop'), exist_ok=True)
os.makedirs(os.path.join(args.output_dir, 'nodetected'), exist_ok=True)
img_files = glob.glob(args.imgs)
for img in img_files:
print(os.path.basename(img))
# test a single image
start = time.time()
result = inference_detector(model, img)
print(time.time() - start)
# show the results
if len(result) < 1 or len(result[0]) < 1:
shutil.copy(img, os.path.join(args.output_dir, 'nodetected'))
continue
output_path = os.path.join(
args.output_dir, 'vis',
os.path.splitext(os.path.basename(img))[0] + ".jpg")
save_result_pyplot(model,
img,
result,
output_path,
score_thr=args.score_thr)
im = Image.open(img).convert("RGB")
for j, bbox in enumerate(result[0]):
if bbox[4] < args.score_thr:
continue
crop_im = im.crop([int(b) for b in bbox[:-1]])
crop_im.save(
os.path.join(
args.output_dir, 'crop',
os.path.splitext(os.path.basename(img))[0] +
"_{}.jpg".format(j)))
def mask_rcnn():
cfg_file = 'configs/mask_rcnn_r50_fpn_1x.py'
ckp_file = '../../checkpoints/' \
'mask_rcnn_r50_fpn_1x_20181010-069fa190.pth'
model = apis.init_detector(cfg_file, ckp_file, device='cuda:0')
img_file = '/home/lhhuang/data/coco/val2017/000000397133.jpg'
img = mmcv.imread(img_file)
result = apis.inference_detector(model, img)
out_file = 'out.jpg'
apis.show_result(
img, result, model.CLASSES,
show=False, out_file=out_file)
def main():
args = parse_args()
# build the model from a config file and a checkpoint file
model = init_detector(args.config_file, args.checkpoint, device='cuda:0')
# test a single image and show the results
img = args.input
result = inference_detector(model, img)
# visualize the results in a new window
# or save the visualization results to image files
show_result(
img, result, model.CLASSES, out_file=img.split('.')[0] + '_result.jpg')
def main(config_name, shrink_coefficient):
config_file = 'local_configs/tianchi/%s.py' % config_name # 修改成自己的配置文件
checkpoint_file = 'work_dirs/%s/latest.pth' % config_name # 修改成自己的训练权重
test_path = '/data/datasets/det/tile_surface_defect_detection/tile_round1_testA_20201231/cut640_imgs/' # 官方测试集图片路径
save_path = 'result/'
if not os.path.exists(save_path):
os.makedirs(save_path)
json_name = os.path.join(
save_path,
'%s_result.json' % time.strftime("%Y%m%d%H%M%S", time.localtime()))
model = init_detector(config_file, checkpoint_file, device='cuda:0')
img_list = []
for img_name in os.listdir(test_path):
if img_name.endswith('.jpg'):
img_list.append(img_name)
result = []
image_num = len(img_list)
for img_id, img_name in enumerate(img_list, 1):
print("%s / %s" % (img_id, image_num))
full_img = os.path.join(test_path, img_name)
predict = inference_detector(model, full_img)
for i, bboxes in enumerate(predict, 1):
if len(bboxes) > 0:
defect_label = i
# print(i)
image_name = img_name
for bbox in bboxes:
x1, y1, x2, y2, score = bbox.tolist()
x1, y1, x2, y2 = round(x1, 2), round(y1, 2), round(
x2, 2), round(y2, 2) #save 0.00
result.append({
'name': image_name,
'category': defect_label,
'bbox': [x1, y1, x2, y2],
'score': score
})
for logs in result:
bbox = logs['bbox']
logs['bbox'] = _shrink_xyxy(bbox, shrink_coefficient)
with open(json_name, 'w') as fp:
json.dump(result,
fp,
separators=(',', ':'),
ensure_ascii=False,
indent=4)
def mask_inference(data_dir, config, ckpt, out_dir):
# build the model from a config file and a checkpoint file
print('Generating first frame mask...')
model = init_detector(config, ckpt, device='cuda:0')
# test a single image
imgs = glob.glob(os.path.join(data_dir, 'JPEGImages/*/*.jpg'))
print('Total images: {}'.format(len(imgs)))
generate_imagesets(data_dir, imgs)
for img in imgs:
# img = '/workspace/solo/test/00001.jpg'
result, cost_time = inference_detector(model, img)
result = filter_result(result, max_num=3)
save_mask(img, result, MASK_THR, out_dir)
def faster_rcnn():
cfg_file = 'configs/faster_rcnn_r50_fpn_1x.py'
ckp_file = '../../checkpoints/' \
'faster_rcnn_r50_fpn_2x_20181010-443129e1.pth'
model = apis.init_detector(cfg_file, ckp_file, device='cuda:0')
img_file = '/home/lhhuang/data/VOCdevkit/VOC2007/JPEGImages/007663.jpg'
img = mmcv.imread(img_file)
result = apis.inference_detector(model, img)
out_file = 'out.jpg'
apis.show_result(
img, result, model.CLASSES,
show=False, out_file=out_file)
def __init__(self,
config_path,
model_path,
verbose_interval=None,
class_restrictions=None):
torch.cuda.empty_cache()
self.model = init_detector(config_path, model_path, device='cuda:0')
# self.class_labels = self.model.CLASSES
self.class_restrictions = class_restrictions
self.verbose_interval = verbose_interval
###
self.i = 0
def init_sess():
PWD = os.path.abspath(os.getcwd())
SDK_PATH = os.path.dirname(os.path.abspath(__file__))
os.chdir(SDK_PATH)
try:
global model, inference_detector
from mmdet.apis import init_detector, inference_detector as do_infer
# build the model from a config file and a checkpoint file
model = init_detector(config_file, checkpoint_file, device='cuda:0')
inference_detector = do_infer
finally:
os.chdir(PWD)
run_sess(np.zeros((480, 640, 3), dtype=np.uint8))
def infer(config,
checkpoint,
img_file_dir,
output_dir,
json_name='bbox_score.json',
show_score_thr=0.3):
model = init_detector(config, checkpoint, device='cuda:0')
img_dir = img_file_dir
file_name_list = os.listdir(img_dir)
img_dir2 = img_dir.replace('_p', '')
results = {}
ik = 0
for i in tqdm(range(len(file_name_list))):
file_name = file_name_list[i]
if os.path.splitext(file_name)[1] not in [
'.jpg', '.png', '.bmp', '.gif'
]:
continue
result_p = inference_detector(model, img_dir + file_name)
result_c = inference_detector(model, img_dir2 + file_name)
if isinstance(result_p, tuple):
bbox_results, _ = result_p
result_p = bbox_results
bbox_results, _ = result_c
result_c = bbox_results
result_above_confidence_num_p = 0
result_above_confidence_num_c = 0
result_p = np.concatenate(result_p)
result_c = np.concatenate(result_c)
for ir in range(len(result_p)):
if result_p[ir, 4] > show_score_thr:
result_above_confidence_num_p = result_above_confidence_num_p + 1
for ir in range(len(result_c)):
if result_c[ir, 4] > show_score_thr:
result_above_confidence_num_c = result_above_confidence_num_c + 1
if result_above_confidence_num_c == 0: # can't find any object in clean img
bb_score = 0
print('i=', ik)
print(file_name)
ik += 1
else:
bb_score = 1 - min(
result_above_confidence_num_c,
result_above_confidence_num_p) / result_above_confidence_num_c
results[file_name] = bb_score
import json
with open(os.path.join(output_dir, json_name), 'w') as f_obj:
json.dump(results, f_obj)
return results
def __init__(self, model_name, model_path):
# make sure these files exist
self.model_name = model_name
self.checkpoint = os.path.join(os.path.dirname(__file__), 'model.pth')
self.config = os.path.join(os.path.dirname(__file__), 'configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py')
self.label_map = parse_classify_rule(os.path.join(os.path.dirname(__file__), 'classify_rule.json'))
self.model = init_detector(self.config, checkpoint=self.checkpoint, device='cpu')
load_checkpoint(self.model, self.checkpoint, map_location='cpu')
self.model.eval()
self.classes = self.model.CLASSES
self.input_image_key = 'images'
self.score = 0.3
def main():
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('--device',
default='cuda:0',
help='Device used for inference')
parser.add_argument('--score-thr',
type=float,
default=0.3,
help='bbox score threshold')
args = parser.parse_args()
hooks = {}
# build the model from a config file and a checkpoint file
model = init_detector(args.config, args.checkpoint, device=args.device)
for name, module in model.named_modules():
print(name)
print(module)
if '1fpn_convs' in name and '.conv' in name:
# if 'backbone.layer1.2.conv3' in name:
# print(name)
# print(module)
hooks[name] = module.register_forward_hook(hook)
# test a single image
result = inference_detector(model, args.img)
for name, module in model.named_modules():
if '2fpn_convs' in name and '.conv' in name:
# if 'backbone.layer1.2.conv3' in name:
out = module._value_hook[0].detach().cpu()
# print(out)
for i in range(128):
plt.box(False)
fig, axarr = plt.subplots()
axarr.axis('off')
axarr.imshow(out[i, :, :], cmap='gray', vmin=0, vmax=1)
# plt.show()
fig.patch.set_visible(False)
fig.savefig('demo/nest94/' + name + str(i) + '.png')
plt.close(fig)
# show the results
# res=filter(result)
show_result_pyplot(model, args.img, result, score_thr=args.score_thr)
def convert2onnx(args, dummy_input):
''' Convert torch model to onnx model '''
# build the model from a config file and a checkpoint file
model = init_detector(args.config, args.checkpoint, device='cuda:0')
if hasattr(model, 'forward_dummy'):
# model.forward = model.extract_feat
model.forward = model.forward_dummy
else:
raise NotImplementedError(
'ONNX exporting is not currently supported with {}'.
format(model.__class__.__name__))
# torch.onnx.export(model, dummy_input, args.out, input_names=['input'], output_names=['outputs'], verbose=True, opset_version=11)
torch.onnx.export(model, dummy_input, args.out, input_names=input_names, output_names=output_names, verbose=True, opset_version=11)
