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)