def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() if os.path.isdir(args.im_or_folder): im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext) else: im_list = [args.im_or_folder] pool = mp.Pool(mp.cpu_count()) func = partial(process_image, model, args) pool.map(func, [im_name for im_name in im_list]) pool.close() pool.join()
def __init__(self): super(Model, self).__init__() self.name = 'Mask RCNN' # Configuration and weights options # By default, we use ResNet50 backbone architecture, you can switch to # ResNet101 to increase quality if your GPU memory is higher than 8GB. # To do so, you will need to download both .yaml and .pkl ResNet101 files # then replace the below 'cfg_file' with the following: # self.cfg_file = 'models/mrcnn/e2e_mask_rcnn_X-101-64x4d-FPN_2x.yaml' self.cfg_file = 'models/mrcnn/e2e_mask_rcnn_R-50-FPN_2x.yaml' self.weights = 'models/mrcnn/model_final.pkl' self.default_cfg = copy.deepcopy(AttrDict(cfg)) # cfg from detectron.core.config self.mrcnn_cfg = AttrDict() self.dummy_coco_dataset = dummy_datasets.get_coco_dataset() # Inference options self.show_box = True self.show_class = True self.thresh = 0.7 self.alpha = 0.4 self.show_border = True self.border_thick = 1 self.bbox_thick = 1 self.font_scale = 0.35 self.binary_masks = False # Define exposed options self.options = ( 'show_box', 'show_class', 'thresh', 'alpha', 'show_border', 'border_thick', 'bbox_thick', 'font_scale', 'binary_masks', ) # Define inputs/outputs self.inputs = {'input': 3} self.outputs = {'output': 3}
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.TEST.WEIGHTS = args.weights cfg.NUM_GPUS = 1 assert_and_infer_cfg() model = infer_engine.initialize_model_from_cfg(args.weights) dummy_dataset = dummy_datasets.get_coco_dataset() if os.path.isdir(args.im_or_folder): im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext) else: im_list = [args.im_or_folder] for i, im_name in enumerate(im_list): logger.info('Processing {}'.format(im_name)) im = cv2.imread(im_name) timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_scores, _, _ = infer_engine.im_detect_all(model, im, None, timers=timers) logger.info('Inference time: {:.3f}s'.format(time.time() - t)) cl = np.argmax(cls_scores) for k, v in timers.items(): logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) logger.info(' | Class is: {}'.format(dummy_dataset.classes[cl])) logger.info(' | Class Confidance is: {:.2f}%'.format( cls_scores[cl] * 100)) if i == 0: logger.info( ' \ Note: inference on the first image will be slower than the ' 'rest (caches and auto-tuning need to warm up)')
def main(args): merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() if os.path.isdir(args.im_or_folder): im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext) else: im_list = [args.im_or_folder] for im_name in im_list: out_name = os.path.join( args.output_dir, '{}'.format(os.path.basename(im_name) + '.' + args.output_ext)) logger.info('Processing {} -> {}'.format(im_name, out_name)) im = cv2.imread(im_name) img1, img2 = image_utils.split_two(im) base_name = os.path.basename(im_name) name_txt = os.path.splitext(base_name)[0] single_process(args, dummy_coco_dataset, img1, base_name + "_1", model) single_process(args, dummy_coco_dataset, img2, base_name + "_2", model)
def get_roidb_and_dataset(dataset_name, proposal_file, ind_range): """Get the roidb for the dataset specified in the global cfg. Optionally restrict it to a range of indices if ind_range is a pair of integers. """ if dataset_name == 'live_targets': from detectron.datasets.live_dataset import LiveRoidb roidb = LiveRoidb() import detectron.datasets.dummy_datasets as dummy_datasets json_dataset = dummy_datasets.get_coco_dataset() if not cfg.TRAIN.USE_FLIPPED: logger.info( 'Live target data set will use flipped examples anyway!') logger.info('"Loaded" dataset: {:s}'.format('live_targets')) return roidb, json_dataset, 0, len(roidb), len(roidb) dataset = JsonDataset(dataset_name) if cfg.TEST.PRECOMPUTED_PROPOSALS: assert proposal_file, 'No proposal file given' roidb = dataset.get_roidb(proposal_file=proposal_file, proposal_limit=cfg.TEST.PROPOSAL_LIMIT) else: roidb = dataset.get_roidb() if ind_range is not None: total_num_images = len(roidb) start, end = ind_range roidb = roidb[start:end] else: start = 0 end = len(roidb) total_num_images = end return roidb, dataset, start, end, total_num_images
def __init__(self, myargv=None, detectron_dir=os.path.expanduser( os.path.join('~', 'libraries', 'detectron'))): workspace.GlobalInit(['caffe2']) setup_logging(__name__) self.logger = logging.getLogger(__name__) args = parse_args(myargv) # config # model_config_filename=os.path.join(detectron_dir, 'configs/12_2017_baselines/e2e_keypoint_rcnn_X-101-32x8d-FPN_s1x.yaml') model_config_filename = os.path.join( detectron_dir, 'configs/12_2017_baselines/e2e_keypoint_rcnn_R-50-FPN_s1x.yaml') args.cfg = model_config_filename merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 # weights # model_weights = 'https://dl.fbaipublicfiles.com/detectron/37732318/12_2017_baselines/e2e_keypoint_rcnn_X-101-32x8d-FPN_s1x.yaml.16_55_09.Lx8H5JVu/output/train/keypoints_coco_2014_train%3Akeypoints_coco_2014_valminusminival/generalized_rcnn/model_final.pkl' model_weights = 'https://dl.fbaipublicfiles.com/detectron/37697714/12_2017_baselines/e2e_keypoint_rcnn_R-50-FPN_s1x.yaml.08_44_03.qrQ0ph6M/output/train/keypoints_coco_2014_train%3Akeypoints_coco_2014_valminusminival/generalized_rcnn/model_final.pkl' cfg.DOWNLOAD_CACHE = os.path.join(detectron_dir, 'detectron-download-cache') args.weights = cache_url(model_weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' self.model = infer_engine.initialize_model_from_cfg(args.weights) self.dummy_coco_dataset = dummy_datasets.get_coco_dataset()
def process_images( args_weights, args_im_or_folder, args_image_ext ): assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' model = infer_engine.initialize_model_from_cfg(args_weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() if os.path.isdir(args_im_or_folder): im_list = glob.iglob(args_im_or_folder + '/*.' + args_image_ext) else: im_list = [args_im_or_folder] for i, im_name in enumerate(im_list): im = cv2.imread(im_name) timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( model, im, None, timers=timers ) boxes, segms, keyps, classes = vis_utils.convert_from_cls_format(cls_boxes, cls_segms, cls_keyps) if classes is not None: class_strs = [dummy_coco_dataset.classes[c] for c in classes] im_name = im_name.split('/')[-1].split('.'+args_image_ext)[0] yield im_name, boxes, segms, class_strs, im.shape[:2]
def __init__( self, cf='/detectron/configs/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml', weights='https://dl.fbaipublicfiles.com/detectron/35861858/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml.02_32_51.SgT4y1cO/output/train/coco_2014_train%3Acoco_2014_valminusminival/generalized_rcnn/model_final.pkl' ): workspace.GlobalInit(['caffe2', '--caffe2_log_level=0']) c2_utils.import_detectron_ops() # OpenCL may be enabled by default in OpenCV3; disable it because it's not # thread safe and causes unwanted GPU memory allocations. cv2.ocl.setUseOpenCL(False) # cfg path self.cfg = cf # weights path self.weights = weights self.weights = cache_url(self.weights, cfg.DOWNLOAD_CACHE) merge_cfg_from_file(self.cfg) cfg.NUM_GPUS = 1 assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' self.model = infer_engine.initialize_model_from_cfg(self.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset()
def main(args): logger = logging.getLogger(__name__) dummy_coco_dataset = dummy_datasets.get_coco_dataset() cfg_orig = load_cfg(envu.yaml_dump(cfg)) im = cv2.imread(args.im_file) if args.rpn_pkl is not None: proposal_boxes, _proposal_scores = get_rpn_box_proposals(im, args) workspace.ResetWorkspace() else: proposal_boxes = None cls_boxes, cls_segms, cls_keyps, cls_bodys = None, None, None, None for i in range(0, len(args.models_to_run), 2): pkl = args.models_to_run[i] yml = args.models_to_run[i + 1] cfg.immutable(False) merge_cfg_from_cfg(cfg_orig) merge_cfg_from_file(yml) if len(pkl) > 0: weights_file = pkl else: weights_file = cfg.TEST.WEIGHTS cfg.NUM_GPUS = 1 assert_and_infer_cfg(cache_urls=False) model = model_engine.initialize_model_from_cfg(weights_file) with c2_utils.NamedCudaScope(0): cls_boxes_, cls_segms_, cls_keyps_ , cls_bodys_= \ model_engine.im_detect_all(model, im, proposal_boxes) cls_boxes = cls_boxes_ if cls_boxes_ is not None else cls_boxes cls_segms = cls_segms_ if cls_segms_ is not None else cls_segms cls_keyps = cls_keyps_ if cls_keyps_ is not None else cls_keyps cls_bodys = cls_bodys_ if cls_bodys_ is not None else cls_bodys workspace.ResetWorkspace() out_name = os.path.join( args.output_dir, '{}'.format(os.path.basename(args.im_file) + '.pdf')) logger.info('Processing {} -> {}'.format(args.im_file, out_name)) with open('test_vis.pkl', 'w') as f: pickle.dump( { 'im': im, 'cls_boxes': np.array(cls_boxes), 'cls_bodys': np.array(cls_bodys) }, f) vis_utils.vis_one_image(im[:, :, ::-1], args.im_file, args.output_dir, cls_boxes, cls_segms, cls_keyps, cls_bodys, dataset=dummy_coco_dataset, box_alpha=0.3, show_class=True, thresh=0.7, kp_thresh=2)
def extract_iuv(image_path, model, infer_engine): im = cv2.imread(image_path) dummy_coco_dataset = dummy_datasets.get_coco_dataset() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps, cls_bodys = infer_engine.im_detect_all( model, im, None ) iuv_out = vis_utils.vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization None, None, cls_boxes, cls_segms, cls_keyps, cls_bodys, dataset=dummy_coco_dataset, box_alpha=0.3, show_class=True, thresh=0.7, kp_thresh=2 ) return iuv_out
def initialization(self): merge_cfg_from_file(self.MODEL_CFG_FILE) cfg.NUM_GPUS = 1 assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' self.model = infer_engine.initialize_model_from_cfg( self.MODEL_WEIGHTS_FILE) self.dummy_coco_dataset = dummy_datasets.get_coco_dataset() meta_data = pd.read_csv(self.THRESHOLDS_FILE).to_dict(orient='records') #print (meta_data) classes = ['background'] thresholds = [1.1] for item in meta_data: name, thr = item['name'], item['threshold'] classes.append(name) thresholds.append(thr) self.dummy_coco_dataset.classes = { i: name for i, name in enumerate(classes) } self.classes = classes self.thresholds = thresholds
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() if os.path.isdir(args.im_or_folder): im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext) else: im_list = [args.im_or_folder] for i, im_name in enumerate(im_list): output_image_dir = os.path.join(args.output_dir, 'images') out_name = os.path.join( output_image_dir, '{}'.format(os.path.basename(im_name) + '.' + args.output_ext)) logger.info('Processing {} -> {}'.format(im_name, out_name)) im = cv2.imread(im_name) timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( model, im, None, timers=timers) logger.info('Inference time: {:.3f}s'.format(time.time() - t)) for k, v in timers.items(): logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) if i == 0: logger.info( ' \ Note: inference on the first image will be slower than the ' 'rest (caches and auto-tuning need to warm up)') # vis_utils.vis_one_image( # im[:, :, ::-1], # BGR -> RGB for visualization # im_name, # output_image_dir, # cls_boxes, # cls_segms, # cls_keyps, # dataset=dummy_coco_dataset, # box_alpha=0.3, # show_class=True, # thresh=0.7, # kp_thresh=2, # ext=args.output_ext, # out_when_no_box=args.out_when_no_box #) _write_to_txt(cls_boxes, cls_segms, cls_keyps, im_name, dummy_coco_dataset)
def __init__(self, cfg_path, weights_path=""): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 weights = cache_url(weights_path, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) self.model = infer_engine.initialize_model_from_cfg(weights) self.dummy_coco_dataset = dummy_datasets.get_coco_dataset()
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 4 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' model = infer_engine.initialize_model_from_cfg(args.weights) # model = infer_engine.initialize_model_from_cfg(args.weights,gpu_id=3) dummy_coco_dataset = dummy_datasets.get_coco_dataset() if os.path.isdir(args.im_or_folder): im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext) else: im_list = [args.im_or_folder] len_str = str(len(os.listdir(args.im_or_folder))) for i, im_name in enumerate(im_list): print('~~~~~~~~~~~~~~~~~' + str(i) + '/' + len_str + "~~~~~~~~~~~~~~~~~~~~~") out_name = os.path.join( args.output_dir, '{}'.format(os.path.basename(im_name) + '.jpg')) logger.info('Processing {} -> {}'.format(im_name, out_name)) im = cv2.imread(im_name) timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( #此处输出的cls_boxes包含类别信息 model, im, None, timers=timers) logger.info('Inference time: {:.3f}s'.format(time.time() - t)) for k, v in timers.items(): logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) if i == 0: logger.info( ' \ Note: inference on the first image will be slower than the ' 'rest (caches and auto-tuning need to warm up)') # ipdb.set_trace() # print(cls_boxes) # print(cls_segms) # print(cls_keyps) vis_utils.vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization im_name, args.output_dir, cls_boxes, cls_segms, cls_keyps, dataset=dummy_coco_dataset, box_alpha=0.3, show_class=True, thresh=0.5, kp_thresh=2, ext='jpg')
def main(args): #logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() cam = cv2.VideoCapture(0) # Set Texture Tex_Atlas = cv2.imread( 'DensePoseData/demo_data/texture_from_SURREAL.png')[:, :, ::-1] / 255.0 #Tex_Atlas = cv2.imread('DensePoseData/demo_data/texture_atlas_200.png')[:,:,::-1] TextureIm = np.zeros([24, 200, 200, 3]) # for i in range(4): for j in range(6): TextureIm[(6 * i + j), :, :, :] = Tex_Atlas[(200 * j):(200 * j + 200), (200 * i):(200 * i + 200), :] while True: if cv2.waitKey(1) & 0xFF == ord('q'): break retval, im = cam.read() #imsmall = cv2.resize(im, None, fx=0.5, fy=0.5) with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps, cls_bodys = infer_engine.im_detect_all( model, im, None, timers=None) iuvout, indsout = vis_utils.vis_webcam(im, cls_boxes, cls_segms, cls_keyps, cls_bodys, thresh=0.9, kp_thresh=2, dataset=dummy_coco_dataset, show_class=True) iuvout, indsout = vis_utils.vis_webcam( im, boxes=cls_boxes, segms=cls_segms, keypoints=cls_keyps, body_uv=cls_bodys, thresh=0.9, kp_thresh=2, dpi=200, box_alpha=0.3, dataset=dummy_coco_dataset, show_class=True, ) #cv2.imshow('input', im) texout = TransferTexturePure(TextureIm, im, iuvout) cv2.imshow('output', texout)
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() submit_result = [] result_file_name = 'detectron_val_result.txt' if os.path.isdir(args.im_or_folder): im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext) else: im_list = [args.im_or_folder] for i, im_name in enumerate(im_list): out_name = os.path.join( args.output_dir, '{}'.format(os.path.basename(im_name) + '.pdf')) logger.info('Processing {} -> {}'.format(im_name, out_name)) im = cv2.imread(im_name) timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( model, im, None, timers=timers) logger.info('Inference time: {:.3f}s'.format(time.time() - t)) for k, v in timers.items(): logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) if i == 0: logger.info( ' \ Note: inference on the first image will be slower than the ' 'rest (caches and auto-tuning need to warm up)') result = vis_utils.vis_one_image_bbox( im[:, :, ::-1], # BGR -> RGB for visualization im_name, args.output_dir, cls_boxes, cls_segms, cls_keyps, dataset=dummy_coco_dataset, box_alpha=0.3, show_class=True, thresh=0.7, kp_thresh=2) if result: submit_result.extend(result) logger.info('Image {}.'.format(i)) # Write file with open(result_file_name, 'wb') as result_file: for item in submit_result: result_file.write("%s\n" % item) logger.info( 'The result file has been written in {}.'.format(result_file_name))
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() if os.path.isdir(args.im_or_folder): im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext) else: im_list = [args.im_or_folder] for i, im_name in enumerate(im_list): out_name = os.path.join( args.output_dir, '{}'.format(os.path.basename(im_name) + '.' + args.output_ext) ) logger.info('Processing {} -> {}'.format(im_name, out_name)) im = cv2.imread(im_name) timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( model, im, None, timers=timers ) logger.info('Inference time: {:.3f}s'.format(time.time() - t)) for k, v in timers.items(): logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) if i == 0: logger.info( ' \ Note: inference on the first image will be slower than the ' 'rest (caches and auto-tuning need to warm up)' ) vis_utils.vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization im_name, args.output_dir, cls_boxes, cls_segms, cls_keyps, dataset=dummy_coco_dataset, box_alpha=0.3, show_class=True, thresh=args.thresh, kp_thresh=args.kp_thresh, ext=args.output_ext, out_when_no_box=args.out_when_no_box )
def GetStats(self, cur_iter, lr): eta_seconds = self.iter_timer.average_time * (cfg.SOLVER.MAX_ITER - cur_iter) eta = str(datetime.timedelta(seconds=int(eta_seconds))) mem_stats = c2_py_utils.GetGPUMemoryUsageStats() mem_usage = np.max(mem_stats['max_by_gpu'][:cfg.NUM_GPUS]) stats = dict( iter=cur_iter, lr=float(lr), time=self.iter_timer.average_time, loss=self.smoothed_total_loss.GetMedianValue(), eta=eta, mb_qsize=int(np.round(self.smoothed_mb_qsize.GetMedianValue())), mem=int(np.ceil(mem_usage / 1024 / 1024)), ) if cfg.TRAIN.DA_FADE_IN: stats['da_weight'] = self.model.da_fade_in.get_weight() if cfg.TRAIN.PADA: stats[ 'avg_pada_weight'] = self.model.class_weight_db.get_avg_pada_weight( ) stats[ 'total_detects'] = self.model.class_weight_db.total_sum_softmax.sum( ) / 2 stats['KL_div'] = self.model.class_weight_db.get_KL_to_init() stats['accuracy_fg'] = self.model.class_weight_db.fg_acc.get() stats[ 'acc_fg_weighted'] = self.model.class_weight_db.weighted_fg_acc.get( ) target_dist = self.model.class_weight_db.get_dist() print('target_dist: {}'.format(list(target_dist))) class_weights = self.model.class_weight_db.class_weights print('class_weights: {}'.format(list(class_weights))) classes = np.array( dummy_datasets.get_coco_dataset().classes.values(), dtype=str) for dist in [target_dist, class_weights]: order = np.argsort(dist)[::-1] o_target_dist = target_dist[order] o_classes = classes[order] cwo = class_weights[order] print("dist tops: ", end='') for prob, w, c in list(zip(o_target_dist, cwo, o_classes))[:5]: print("{}:{:.3f} ({:.3f})".format(c, prob, w), end='; ') print() print() for k, v in self.smoothed_losses_and_metrics.items(): stats[k] = v.GetMedianValue() # for k,v in stats.items(): # print(k,v) return stats
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() ''' add onboard cam support ''' cam = cv2.VideoCapture("nvcamerasrc ! video/x-raw(memory:NVMM), width=(int)1280, height=(int)720,format=(string)I420, framerate=(fraction)30/1 ! nvvidconv flip-method=0 ! video/x-raw, format=(string)BGRx ! videoconvert ! video/x-raw, format=(string)BGR ! appsink") if cam.isOpened(): print("camara open succeded") im_name = "Detection" while True: _, im = cam.read(); timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( model, im, None, timers=timers ) logger.info('Inference time: {:.3f}s'.format(time.time() - t)) ''' for k, v in timers.items(): logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) ''' vis_utils.vis_cam_image( im[:, :, ::-1], # BGR -> RGB for visualization im_name, cls_boxes, cls_segms, cls_keyps, dataset=dummy_coco_dataset, box_alpha=0.3, show_class=True, thresh=0.7, kp_thresh=2 ) """ key=cv2.waitKey(10) if key == 27: # Check for ESC key cv2.destroyAllWindows() break ; """ else: print("camera open failed")
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() logger.info( ' \ Note: inference on the first image will be slower than the ' 'rest (caches and auto-tuning need to warm up)' ) while True: reply = {} # Wait for next request from client im,extra = zmqa.recv(socket) if extra is not None and 'fname' in extra: print("Received request %s" % extra) reply['fname']=extra['fname'] else: print("Received request %s" % extra) reply['fname'] = 'input.jpg' # set file name im_name=reply['fname'] timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( model, im, None, timers=timers ) logger.info('Inference time: {:.3f}s'.format(time.time() - t)) for k, v in timers.items(): logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) out_name = os.path.join( args.output_dir, '{}'.format(os.path.basename(im_name) + '.' + args.output_ext) ) masks=do_one_image_opencv( im, # BGR -> RGB for visualization cls_boxes, cls_segms, cls_keyps, thresh=0.7, reply=reply ) zmqa.send(socket, masks,extra=reply)
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() # load match scene match_scene_dict = pickle.load(open('{}/pkl/match_scene_intervals_dict.pkl'.format(args.data_dir), 'rb')) HW_foreground = set() JZ_foreground = set() for match in match_scene_dict['HW_foreground']: for fid in range(match[1], match[2]): HW_foreground.add(fid) for match in match_scene_dict['JZ_foreground']: for fid in range(match[1], match[2]): JZ_foreground.add(fid) # detect frame cap = cv2.VideoCapture('{}/videos/{}'.format(args.data_dir, args.input_video)) result_list = [] fid = 0 ret = True # create dir for player in ['HW', 'JZ']: dir = '{data_dir}/image/img_{player}'.format(data_dir=args.data_dir, player=player) if not os.path.exists(dir): os.makedirs(dir) dir = '{data_dir}/image/densepose_{player}'.format(data_dir=args.data_dir, player=player) if not os.path.exists(dir): os.makedirs(dir) vid = 65 while ret: if fid % 1000 == 0: print("Inferring frame %d" % fid) ret, image = cap.read() if fid in HW_foreground or fid in JZ_foreground: player = 'HW' if fid in HW_foreground else 'JZ' img_path = '{data_dir}/image/img_{player}/img_{vid}_{fid}_{player}.jpg'.format(data_dir=args.data_dir, player=player, vid=vid, fid=fid) pose_path = '{data_dir}/image/densepose_{player}/densepose_{vid}_{fid}_{player}.jpg'.format(data_dir=args.data_dir, player=player, vid=vid, fid=fid) result_one_image = infer_one_frame(image, model, img_path, pose_path) result_one_image['foreground'] = player result_one_image['fid'] = fid result_list.append(result_one_image) fid += 1 # if len(result_list) > 100: # break pickle.dump(result_list, open('{}/pkl/result.pkl'.format(args.data_dir), 'wb'), protocol=2)
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() if os.path.isdir(args.im_or_folder): im_list = glob.iglob(args.im_or_folder + '/*/*.' + args.image_ext) else: im_list = [args.im_or_folder] for i, im_name in enumerate(im_list): vid = im_name.split("/")[1] output_dir = args.output_dir + vid try: os.mkdir(output_dir) except OSError: pass out_name = os.path.join( output_dir, '{}'.format(os.path.basename(im_name) + '.pdf')) if os.path.exists(out_name): continue logger.info('Processing {} -> {}'.format(im_name, out_name)) im = cv2.imread(im_name) timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps, cls_bodys = infer_engine.im_detect_all( model, im, None, timers=timers) logger.info('Inference time: {:.3f}s'.format(time.time() - t)) for k, v in timers.items(): logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) if i == 0: logger.info( ' \ Note: inference on the first image will be slower than the ' 'rest (caches and auto-tuning need to warm up)') vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization im_name, output_dir, cls_boxes, cls_segms, cls_keyps, cls_bodys, dataset=dummy_coco_dataset, box_alpha=0.3, show_class=True, thresh=0.7, kp_thresh=2)
def main(args): #ros initialization rospy.init_node('get_image', anonymous=True) args.im_or_folder = get_image() while not rospy.is_shutdown(): merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() # set im list to be the ros image im = args.im_or_folder timers = defaultdict(Timer) with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( model, im, None, timers=timers) # calls the method that performs actual detection and inference fig = vis_utils.vis_one_image_opencv( im[:, :, ::-1], # BGR -> RGB for visualization cls_boxes, cls_segms, cls_keyps, dataset=dummy_coco_dataset, show_class=True, thresh=args.thresh, kp_thresh=args.kp_thresh, ) # img is rgb, convert to opencv's default bgr img = cv2.cvtColor(fig, cv2.COLOR_RGB2BGR) image_publisher = rospy.Publisher('detectron_output', Image, queue_size=10) brdg = CvBridge() image_publisher.publish(brdg.cv2_to_imgmsg(img, "bgr8")) try: start_time = time.time() print(time.time() - start_time) main(args) except KeyboardInterrupt: print("shutting down") cv2.destroyAllWindows()
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() print("capturing video " + args.input) cap = cv2.VideoCapture(args.input) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # pdb.set_trace() grab = 1 if (cap.isOpened() == False): print("Error opening video stream or file") exit while (cap.isOpened() and grab <= total_frames): grab += 1 ret_val, im = cap.read() #skips intermediate frames #if grab%2 !=0: # continue #uncomment to resize image #im = cv2.resize(im, (int(1280/1),int(720/1))) timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps, cls_bodys = infer_engine.im_detect_all( model, im, None, timers=timers) output_name = 'frame' + str(grab).zfill(4) + '.mp4' print("| Analysed frame {0} / {1} in {2}ms".format( grab, total_frames, int(1000. * (time.time() - t)))) #print('\t | Inference time: {:.3f}s'.format(time.time() - t)) #for k, v in timers.items(): # print('\t | {}: {:.3f}s'.format(k, v.average_time)) ret = vis_utils.vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization output_name, args.output_dir, cls_boxes, cls_segms, cls_keyps, cls_bodys, dataset=dummy_coco_dataset, box_alpha=0.3, show_class=False, thresh=0.7, kp_thresh=2) cap.release() cv2.destroyAllWindows()
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() h5py_file_path = ('/media/hdd1/tanya/open-pose/' 'paired_filenames512_image_keypoints512_main.h5') hf = h5py.File(h5py_file_path, 'r') IUV_image_list = [] for i in range(len(hf['images'])): im = hf['images'][i] timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps, cls_bodys = infer_engine.im_detect_all( model, im, None, timers=timers ) logger.info('Inference time: {:.3f}s'.format(time.time() - t)) for k, v in timers.items(): logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) if i == 0: logger.info( ' \ Note: inference on the first image will be slower than the ' 'rest (caches and auto-tuning need to warm up)' ) IUV_image = vis_utils.vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization 'dummy', args.output_dir, cls_boxes, cls_segms, cls_keyps, cls_bodys, dataset=dummy_coco_dataset, box_alpha=0.3, show_class=True, thresh=0.7, kp_thresh=2) IUV_image_list.append(IUV_image) IUV_images_final = np.stack(IUV_image_list, 0) with h5py.File('./paired_filenames512_image_keypoints_withIUV.h5', 'w') as f: f.create_dataset('images', data=np.array(hf['images'])) f.create_dataset('keypoints', data=np.array(hf['keypoints'])) f.create_dataset('IUV', data=IUV_images_final)
def get_detectron_result(im): """ 功能: 获取传入图像的检测结果 输入参数列表: im: BGR格式图片 返回参数列表: img: 检测结果图片 detectron_result_info: 检测结果字典 """ setup_logging(__name__) logger = logging.getLogger(__name__) dummy_coco_dataset = dummy_datasets.get_coco_dataset() cfg_orig = load_cfg(yaml.dump(cfg)) if rpn_pkl is not None: proposal_boxes, _proposal_scores = get_rpn_box_proposals(im) workspace.ResetWorkspace() else: proposal_boxes = None cls_boxes, cls_segms, cls_keyps = None, None, None pkl = rpn_pkl yml = rpn_cfg cfg.immutable(False) merge_cfg_from_cfg(cfg_orig) merge_cfg_from_file(yml) if len(pkl) > 0: weights_file = pkl else: weights_file = cfg.TEST.WEIGHTS cfg.NUM_GPUS = 1 assert_and_infer_cfg(cache_urls=False) model = model_engine.initialize_model_from_cfg(weights_file) with c2_utils.NamedCudaScope(0): cls_boxes_, cls_segms_, cls_keyps_ = \ model_engine.im_detect_all(model, im, proposal_boxes) cls_boxes = cls_boxes_ if cls_boxes_ is not None else cls_boxes cls_segms = cls_segms_ if cls_segms_ is not None else cls_segms cls_keyps = cls_keyps_ if cls_keyps_ is not None else cls_keyps workspace.ResetWorkspace() img, detectron_result_info = vis_utils.vis_one_image_opencv( im, cls_boxes, cls_segms, cls_keyps, dataset=dummy_coco_dataset, show_box=True, show_class=True) return img, detectron_result_info
def main(args): logger = logging.getLogger(__name__) dummy_coco_dataset = dummy_datasets.get_coco_dataset() cfg_orig = load_cfg(yaml.dump(cfg)) im = cv2.imread(args.im_file) if args.rpn_pkl is not None: proposal_boxes, _proposal_scores = get_rpn_box_proposals(im, args) workspace.ResetWorkspace() else: proposal_boxes = None cls_boxes, cls_segms, cls_keyps = None, None, None for i in range(0, len(args.models_to_run), 2): pkl = args.models_to_run[i] yml = args.models_to_run[i + 1] cfg.immutable(False) merge_cfg_from_cfg(cfg_orig) merge_cfg_from_file(yml) if len(pkl) > 0: weights_file = pkl else: weights_file = cfg.TEST.WEIGHTS cfg.NUM_GPUS = 1 assert_and_infer_cfg(cache_urls=False) model = model_engine.initialize_model_from_cfg(weights_file) with c2_utils.NamedCudaScope(0): cls_boxes_, cls_segms_, cls_keyps_ = \ model_engine.im_detect_all(model, im, proposal_boxes) cls_boxes = cls_boxes_ if cls_boxes_ is not None else cls_boxes cls_segms = cls_segms_ if cls_segms_ is not None else cls_segms cls_keyps = cls_keyps_ if cls_keyps_ is not None else cls_keyps workspace.ResetWorkspace() out_name = os.path.join( args.output_dir, '{}'.format(os.path.basename(args.im_file) + '.pdf') ) logger.info('Processing {} -> {}'.format(args.im_file, out_name)) vis_utils.vis_one_image( im[:, :, ::-1], args.im_file, args.output_dir, cls_boxes, cls_segms, cls_keyps, dataset=dummy_coco_dataset, box_alpha=0.3, show_class=True, thresh=0.7, kp_thresh=2 )
def run_task(clip, background_image): workspace.GlobalInit(['caffe2', '--caffe2_log_level=0']) merge_cfg_from_file('/smartcrop/configs/e2e_mask_rcnn_R-101-FPN_2x.yaml') assert_and_infer_cfg() model = infer_engine.initialize_model_from_cfg( '/smartcrop/weights/e2e_mask_rcnn_R-101-FPN_2x/model_final.pkl') dataset = dummy_datasets.get_coco_dataset() clip = clip.fl_image( lambda image: extract_person(image, background_image, model, dataset)) return clip
def setup(self): c2_utils.import_detectron_ops() cv2.ocl.setUseOpenCL(False) merge_cfg_from_file( '/DensePose/configs/DensePose_ResNet101_FPN_s1x-e2e.yaml') cfg.NUM_GPUS = 1 weights = cache_url( 'https://s3.amazonaws.com/densepose/DensePose_ResNet101_FPN_s1x-e2e.pkl', cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) model = infer_engine.initialize_model_from_cfg(weights, 1) dummy_coco_dataset = dummy_datasets.get_coco_dataset() return model
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 4 args.weights = '/detectron-download-cache/35861858/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml.02_32_51.SgT4y1cO/output/train/coco_2014_train:coco_2014_valminusminival/generalized_rcnn/model_final.pkl' #cache_url(args.weights, cfg.DOWNLOAD_CACHE) print (args.weights) assert_and_infer_cfg(cache_urls=False) model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() if os.path.isdir(args.im_or_folder): im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext) else: im_list = [args.im_or_folder] cv2.namedWindow("Video", cv2.WND_PROP_FULLSCREEN) cv2.setWindowProperty("Video",cv2.WND_PROP_FULLSCREEN,cv2.WINDOW_FULLSCREEN) cap=cv2.VideoCapture(0) width = cap.get(3) height = cap.get(4) cap.set(3,1280) cap.set(4,720) print (width, height) while True: ret, im = cap.read() timers = defaultdict(Timer) im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( model, im, None, timers=timers ) deframe = vis_utils.vis_one_image_opencv( im[:, :, ::-1], # BGR -> RGB for visualization cls_boxes, cls_segms, cls_keyps, dataset=dummy_coco_dataset, show_class = True, thresh=0.7, kp_thresh=2, show_box=True ) cv2.imshow('Video',deframe) if cv2.waitKey(1) & 0xFF == ord('q'): break cap.release() cv2.destroyAllWindows()
def main(args): logger = logging.getLogger(__name__) pdb.set_trace() merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() if os.path.isdir(args.im_or_folder): im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext) else: im_list = [args.im_or_folder] print(im_list) for i, im_name in enumerate(im_list): im_basename = os.path.basename(im_name) out_name = os.path.join( args.output_dir, '{}'.format(im_basename + '.' + args.output_ext)) logger.info('Processing {} -> {}'.format(im_name, out_name)) im = cv2.imread(im_name) timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( model, im, None, timers=timers) result_file = os.path.join(args.output_dir, im_basename + '_bbox.pkl') with open(result_file, 'wb') as f: boxes, _, _, _ = vis_utils.convert_from_cls_format( cls_boxes, cls_segms, cls_keyps) print(boxes.shape) pickle.dump(cls_boxes, f) logger.info('Inference time: {:.3f}s'.format(time.time() - t)) for k, v in timers.items(): logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) if i == 0: logger.info( ' \ Note: inference on the first image will be slower than the ' 'rest (caches and auto-tuning need to warm up)')
def main(args): logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() zmq_video = args.source == "zmq" frameId = 0 if zmq_video: context = zmq.Context() socket = context.socket(zmq.REP) socket.bind("tcp://*:7001") else: # From virtual camera video and its associated timestamp file on Drive PX2,e.g."./lane/videofilepath.h264" cap = cv2.VideoCapture(args.source) while True: if zmq_video: try: message = socket.recv() print("Received message length:" + str(len(message)) + " type:" + str(type(message))) socket.send("ok") position = message.find(ZMQ_SEPER, 0, 100) frameId = message[:position] message = message[position + len(ZMQ_SEPER):] img_np = np.fromstring(message, np.uint8) img_np = img_np.reshape((400, 1400, 3)) print("nparr type:" + str(type(img_np)) + " shape:" + str(img_np.shape)) ret = True except KeyboardInterrupt: print("interrupt received, stopping...") socket.close() context.term() ret = False cap.release() else: ret, img_np = cap.read() frameId += 1 # read completely or raise exception if not ret: print("cannot get frame") break hanle_frame(args, frameId, img_np, logger, model, dummy_coco_dataset)
def main(args): glob_keypoints = [] logger = logging.getLogger(__name__) merge_cfg_from_file(args.cfg) cfg.NUM_GPUS = 1 args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) assert_and_infer_cfg(cache_urls=False) assert not cfg.MODEL.RPN_ONLY, \ 'RPN models are not supported' assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ 'Models that require precomputed proposals are not supported' model = infer_engine.initialize_model_from_cfg(args.weights) dummy_coco_dataset = dummy_datasets.get_coco_dataset() if os.path.isdir(args.im_or_folder): im_list = glob.iglob(args.im_or_folder + '/*' + '.png') else: im_list = [args.im_or_folder] im_list = sorted(im_list) for i, im_name in enumerate(im_list): out_name = os.path.join( args.output_dir, '{}'.format(os.path.basename(im_name) + '.' + args.output_ext) ) logger.info('Processing {} -> {}'.format(im_name, out_name)) im = cv2.imread(im_name) timers = defaultdict(Timer) t = time.time() with c2_utils.NamedCudaScope(0): cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( model, im, None, timers=timers ) logger.info('Inference time: {:.3f}s'.format(time.time() - t)) for k, v in timers.items(): logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) if i == 0: logger.info( ' \ Note: inference on the first image will be slower than the ' 'rest (caches and auto-tuning need to warm up)' ) vis_utils.vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization im_name, args.output_dir, cls_boxes, cls_segms, cls_keyps, dataset=dummy_coco_dataset, box_alpha=0.3, show_class=True, thresh=args.thresh, kp_thresh=args.kp_thresh, ext=args.output_ext, out_when_no_box=args.out_when_no_box ) cls_boxes_np = np.asarray(cls_boxes) cls_boxes_prob = cls_boxes_np[1][:,4] idx_max_prob = np.argmax(cls_boxes_prob) cls_keyps_max_prob = cls_keyps[1][idx_max_prob] pose_x_y_prob_after_softmax = cls_keyps_max_prob[[0,1,3]] glob_keypoints.append(np.transpose(pose_x_y_prob_after_softmax)) dictionarry_keypoints={'S1': {'Directions 1' : np.asarray([glob_keypoints])}} metadata = {'layout_name': 'h36m', 'num_joints': 17, 'keypoints_symmetry': [[4, 5, 6, 11, 12, 13], [1, 2, 3, 14, 15, 16]]} #np.savez(os.path.join('/home/narvis/Dev/VideoPose3D/data', "data_2d_detections.npz"), metadata=metadata, positions_2d=dictionarry_keypoints) np.savez(os.path.join(args.output_dir, "data_2d_detections.npz"), metadata=metadata, positions_2d=dictionarry_keypoints)