def main(args): cfg = setup(args) model = DefaultTrainer.build_model(cfg) modules = model.get_conv_bn_modules() for m in modules: print(m.weight.size()) load_darknet_weights(args.initial_weights, modules) save_path = os.path.join(args.output_dir, "yolov3.pth") torch.save(model.state_dict(), save_path) print("model save to", save_path)
def __init__(self, config_file, mode='caffe', extract_mode='roi_feats', min_max_boxes='10,50'): args = {} args['config_file'] = config_file args['mode'] = mode args['extract_mode'] = extract_mode args['min_max_boxes'] = min_max_boxes args['eval_only'] = True args = Pack(args) cfg = setup(args) self.model = DefaultTrainer.build_model(cfg) DetectionCheckpointer(self.model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=False ) self.model.eval() self.cfg = cfg self.vg_objects = load_vocab_file('objects_vocab.txt')
def evaluate(dataset): register_one_set(dataset) cfg = get_my_cfg() cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 model = DefaultTrainer.build_model( cfg) # just built the model without weights checkpoiner = DetectionCheckpointer(model, cfg.OUTPUT_DIR) checkpoiner.resume_or_load( cfg.MODEL.WEIGHTS, resume=False) # loaded the weights we had trained evaluator = COCOEvaluator(dataset, ("bbox", ), False, output_dir=os.path.join("output", "evaluate")) loader = build_detection_test_loader(cfg, dataset) print(inference_on_dataset(model, loader, evaluator))
def main(args): # if args.unitest: # return unitest() cfg = setup(args) for d in ["train", 'val']: # train for 6998images , val for 1199 images DatasetCatalog.register("chefCap_" + d, lambda d=d: get_chefcap_image_dicts()) MetadataCatalog.get("chefCap_" + d).set( thing_classes=list(things_class_dict.keys())) if d == 'val': MetadataCatalog.get("chefCap_val").evaluator_type = "pascal_voc" MetadataCatalog.get("chefCap_val").year = 2012 MetadataCatalog.get( "chefCap_val" ).dirname = "/opt/work/chefCap/detectron2_fasterrcnn/data" # for d in ["/opt/work/chefCap/data/ziped/Making-PascalVOC-export/"]: # DatasetCatalog.register("chefCap_val", # lambda d=d: get_chefcap_image_dicts(d)) # MetadataCatalog.get("chefCap_val").set( # thing_classes=['face-head', 'mask-head', 'face-cap', 'mask-cap']) # MetadataCatalog.get("chefCap_val").evaluator_type = "pascal_voc" # MetadataCatalog.get("chefCap_val").dirname = "/opt/work/chefCap/data/ziped/Making-PascalVOC-export/" # MetadataCatalog.get("chefCap_val").year = 2012 if args.eval_only: model = DefaultTrainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) res = DefaultTrainer.test(cfg, model) if cfg.TEST.AUG.ENABLED: res.update(DefaultTrainer.test_with_TTA(cfg, model)) if comm.is_main_process(): verify_results(cfg, res) return res trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=args.resume) # if cfg.TEST.AUG.ENABLED: # trainer.register_hooks( # [hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))] # ) return trainer.train()
def load_model_img(): global net_img global cfg time_start = time.time() print('\nLoad image feature extraction model') logging.disable() parser = argparse.ArgumentParser() parser.add_argument( '--config-file', default='model/extract-bua-caffe-r101.yaml', metavar='FILE', help='path to config file', ) parser.add_argument('--mode', default='caffe', type=str, help='bua_caffe, ...') parser.add_argument( '--resume', action='store_true', help='whether to attempt to resume from the checkpoint directory', ) parser.add_argument( 'opts', help='Modify config options using the command-line', default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args([]) cfg = setup(args) MIN_BOXES = 10 MAX_BOXES = 100 CONF_THRESH = 0.2 net_img = DefaultTrainer.build_model(cfg) DetectionCheckpointer(net_img, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) net_img.eval() time_end = time.time() print('Finish load net model!') print('Model load time: {:.3f}s\n'.format(time_end - time_start))
def __init__(self, min_boxes=3, max_boxes=10, threshold=0.5): config_file = 'bottom-up-attention.pytorch/configs/bua-caffe/extract-bua-caffe-r101.yaml' self._cfg = get_cfg() add_bottom_up_attention_config(self._cfg, True) self._cfg.merge_from_file(config_file) self._cfg.MODEL.DEVICE = 'cpu' self._model = DefaultTrainer.build_model(self._cfg) DetectionCheckpointer(self._model, save_dir=self._cfg.OUTPUT_DIR).resume_or_load(self._cfg.MODEL.WEIGHTS) self._model.eval() self._min_boxes = min_boxes self._max_boxes = max_boxes self._threshold = threshold self._classes = ['__background__'] with open(os.path.join('bottom-up-attention.pytorch', 'evaluation', 'objects_vocab.txt')) as f: for object in f.readlines(): self._classes.append(object.split(',')[0].lower().strip())
def main(args): cfg = setup(args) show = True register_openlogo(cfg.DATASETS.TRAIN[0], "datasets/data/openlogo", "trainval", "supervised_imageset") register_openlogo(cfg.DATASETS.TEST[0], "datasets/data/openlogo", "test", "supervised_imageset") trainer = DefaultTrainer(cfg) evaluator = OpenLogoDetectionEvaluator(cfg.DATASETS.TEST[0]) if args.eval_only: model = trainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) if show: visualize(cfg, amount=20) res = trainer.test(cfg, model, evaluators=[evaluator]) if comm.is_main_process(): verify_results(cfg, res) if cfg.TEST.AUG.ENABLED: res.update(trainer.test_with_TTA(cfg, model)) return res trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=args.resume) if cfg.TEST.AUG.ENABLED: trainer.register_hooks([ hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model)) ]) return trainer.train()
def main(args): cfg = setup(args) model = DefaultTrainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume ) save_onnx_name = cfg.MODEL.WEIGHTS.replace('.pth', '.onnx') model.eval() inputs = torch.randn(1, 3, cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST).to(model.device) with torch.no_grad(): torch.onnx.export( model, inputs, save_onnx_name, export_params=True, verbose=False, training=False, do_constant_folding=True, input_names=['input'], # output_names=['output'], # opset_version=9, )
cfg.DATASETS.TEST = ("openimages_val", ) cfg.DATASETS.TRAIN = ("openimages_train", ) predictor = DefaultPredictor(cfg) # trainer = DefaultTrainer(cfg) if args.mode == 'infer': val_dicts = get_val_dicts() for d in random.sample(val_dicts, 10): im = cv2.imread(d['file_name']) outputs = predictor(im) v = Visualizer(im[:, :, ::-1], metadata=openimages_val_metadata, scale=0.8, instance_mode=ColorMode.IMAGE_BW) v = v.draw_instance_predictions(outputs['instances'].to('cpu')) cv2.imshow('image', v.get_image()[:, :, ::-1]) cv2.waitKey(0) if args.mode == 'evaluate': print(cfg.DATASETS.TEST[0]) trainer = DefaultTrainer(cfg) model = trainer.build_model(cfg) DetectionCheckpointer(model).load(args.model_pth) evaluator = COCOEvaluator(cfg.DATASETS.TEST[0], cfg, False, output_dir="./output/") val_loader = build_detection_test_loader(cfg, "openimages_val") inference_on_dataset(model, val_loader, evaluator)
def build_tsv(gpu_id: int = 0): args = Argument() cfg = setup(args) # os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id model = DefaultTrainer.build_model(cfg) DetectionCheckpointer(model).resume_or_load( cfg.MODEL.WEIGHTS, resume=False, ) model.eval() # Set up the simulator print("Starting sim") sim = MatterSim.Simulator() sim.setCameraResolution(WIDTH, HEIGHT) sim.setCameraVFOV(math.radians(VFOV)) sim.setDiscretizedViewingAngles(False) sim.setBatchSize(1) sim.setPreloadingEnabled(False) sim.initialize() print("Sim has start") # import pdb # pdb.set_trace() print("STARTING") with open(OUTFILE, "wt") as tsvfile: writer = csv.DictWriter(tsvfile, delimiter="\t", fieldnames=TSV_FIELDNAMES) # Loop all the viewpoints in the simulator print("loading viewpoints") viewpoint_ids = load_viewpointids() for scanId, viewpointId in tqdm(viewpoint_ids): # Loop all discretized views from this location ims = [] sim.newEpisode( [scanId], [viewpointId], [0], [math.radians(ELEVATION_START)] ) for ix in range(VIEWPOINT_SIZE): state = sim.getState()[0] # Transform and save generated image ims.append(transform_img(state.rgb)) # Build state if ix == 0: record = { "scanId": state.scanId, "viewpointId": state.location.viewpointId, "viewHeading": np.zeros(VIEWPOINT_SIZE, dtype=np.float32), "viewElevation": np.zeros(VIEWPOINT_SIZE, dtype=np.float32), "image_h": HEIGHT, "image_w": WIDTH, "vfov": VFOV, } record["viewHeading"][ix] = state.heading record["viewElevation"][ix] = state.elevation # Move the sim viewpoint so it ends in the same place elev = 0.0 heading_chg = math.pi * 2 / VIEWS_PER_SWEEP view = ix % VIEWS_PER_SWEEP sweep = ix // VIEWS_PER_SWEEP if view + 1 == VIEWS_PER_SWEEP: # Last viewpoint in sweep elev = math.radians(ELEVATION_INC) sim.makeAction([0], [heading_chg], [elev]) # Run detection for ix in range(VIEWPOINT_SIZE): get_detections_from_im(record, model, ims[ix], cfg) filter(record, MAX_TOTAL_BOXES) for k, v in record.items(): if isinstance(v, np.ndarray): record[k] = str(base64.b64encode(v), "utf-8") writer.writerow(record)
def extract_feat(split_idx, img_list, cfg, args, actor: ActorHandle): num_images = len(img_list) print('Number of images on split{}: {}.'.format(split_idx, num_images)) model = DefaultTrainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) model.eval() generate_npz_list = [] for im_file in (img_list): if os.path.exists( os.path.join(args.output_dir, im_file.split('.')[0] + '.npz')): actor.update.remote(1) continue im = cv2.imread(os.path.join(args.image_dir, im_file)) if im is None: print(os.path.join(args.image_dir, im_file), "is illegal!") actor.update.remote(1) continue dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN) # extract roi features if cfg.MODEL.BUA.EXTRACTOR.MODE == 1: attr_scores = None with torch.set_grad_enabled(False): if cfg.MODEL.BUA.ATTRIBUTE_ON: boxes, scores, features_pooled, attr_scores = model( [dataset_dict]) else: boxes, scores, features_pooled = model([dataset_dict]) boxes = [box.tensor.cpu() for box in boxes] scores = [score.cpu() for score in scores] features_pooled = [feat.cpu() for feat in features_pooled] if not attr_scores is None: attr_scores = [attr_score.cpu() for attr_score in attr_scores] generate_npz_list.append( generate_npz.remote(1, actor, args, cfg, im_file, im, dataset_dict, boxes, scores, features_pooled, attr_scores)) # extract bbox only elif cfg.MODEL.BUA.EXTRACTOR.MODE == 2: with torch.set_grad_enabled(False): boxes, scores = model([dataset_dict]) boxes = [box.cpu() for box in boxes] scores = [score.cpu() for score in scores] generate_npz_list.append( generate_npz.remote(2, actor, args, cfg, im_file, im, dataset_dict, boxes, scores)) # extract roi features by bbox elif cfg.MODEL.BUA.EXTRACTOR.MODE == 3: if not os.path.exists( os.path.join(args.bbox_dir, im_file.split('.')[0] + '.npz')): actor.update.remote(1) continue bbox = torch.from_numpy( np.load( os.path.join(args.bbox_dir, im_file.split('.')[0] + '.npz'))['bbox']) * dataset_dict['im_scale'] proposals = Instances(dataset_dict['image'].shape[-2:]) proposals.proposal_boxes = BUABoxes(bbox) dataset_dict['proposals'] = proposals attr_scores = None with torch.set_grad_enabled(False): if cfg.MODEL.BUA.ATTRIBUTE_ON: boxes, scores, features_pooled, attr_scores = model( [dataset_dict]) else: boxes, scores, features_pooled = model([dataset_dict]) boxes = [box.tensor.cpu() for box in boxes] scores = [score.cpu() for score in scores] features_pooled = [feat.cpu() for feat in features_pooled] if not attr_scores is None: attr_scores = [ attr_score.data.cpu() for attr_score in attr_scores ] generate_npz_list.append( generate_npz.remote(3, actor, args, cfg, im_file, im, dataset_dict, boxes, scores, features_pooled, attr_scores)) ray.get(generate_npz_list)
def main(): parser = argparse.ArgumentParser( description="PyTorch Object Detection2 Inference") parser.add_argument( "--config-file", default="configs/bua-caffe/extract-bua-caffe-r101.yaml", metavar="FILE", help="path to config file", ) parser.add_argument("--mode", default="caffe", type=str, help="bua_caffe, ...") parser.add_argument('--out-dir', dest='output_dir', help='output directory for features', default="features") parser.add_argument('--image-dir', dest='image_dir', help='directory with images', default="image") parser.add_argument( "--resume", action="store_true", help="whether to attempt to resume from the checkpoint directory", ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() cfg = setup(args) MIN_BOXES = 10 MAX_BOXES = 100 CONF_THRESH = 0.2 model = DefaultTrainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) # Extract features. imglist = os.listdir(args.image_dir) num_images = len(imglist) print('Number of images: {}.'.format(num_images)) model.eval() for im_file in tqdm.tqdm(imglist): im = cv2.imread(os.path.join(args.image_dir, im_file)) dataset_dict = get_image_blob(im) with torch.set_grad_enabled(False): # boxes, scores, features_pooled = model([dataset_dict]) if cfg.MODEL.BUA.ATTRIBUTE_ON: boxes, scores, features_pooled, attr_scores = model( [dataset_dict]) else: boxes, scores, features_pooled = model([dataset_dict]) dets = boxes[0].tensor.cpu() / dataset_dict['im_scale'] scores = scores[0].cpu() feats = features_pooled[0].cpu() max_conf = torch.zeros((scores.shape[0])).to(scores.device) for cls_ind in range(1, scores.shape[1]): cls_scores = scores[:, cls_ind] keep = nms(dets, cls_scores, 0.3) max_conf[keep] = torch.where(cls_scores[keep] > max_conf[keep], cls_scores[keep], max_conf[keep]) keep_boxes = torch.nonzero(max_conf >= CONF_THRESH).flatten() if len(keep_boxes) < MIN_BOXES: keep_boxes = torch.argsort(max_conf, descending=True)[:MIN_BOXES] elif len(keep_boxes) > MAX_BOXES: keep_boxes = torch.argsort(max_conf, descending=True)[:MAX_BOXES] image_feat = feats[keep_boxes] image_bboxes = dets[keep_boxes] image_objects_conf = np.max(scores[keep_boxes].numpy(), axis=1) image_objects = np.argmax(scores[keep_boxes].numpy(), axis=1) if cfg.MODEL.BUA.ATTRIBUTE_ON: attr_scores = attr_scores[0].cpu() image_attrs_conf = np.max(attr_scores[keep_boxes].numpy(), axis=1) image_attrs = np.argmax(attr_scores[keep_boxes].numpy(), axis=1) info = { 'image_id': im_file.split('.')[0], 'image_h': np.size(im, 0), 'image_w': np.size(im, 1), 'num_boxes': len(keep_boxes), 'objects_id': image_objects, 'objects_conf': image_objects_conf, 'attrs_id': image_attrs, 'attrs_conf': image_attrs_conf, } else: info = { 'image_id': im_file.split('.')[0], 'image_h': np.size(im, 0), 'image_w': np.size(im, 1), 'num_boxes': len(keep_boxes), 'objects_id': image_objects, 'objects_conf': image_objects_conf } output_file = os.path.join(args.output_dir, im_file.split('.')[0]) np.savez_compressed(output_file, x=image_feat, bbox=image_bboxes, num_bbox=len(keep_boxes), image_h=np.size(im, 0), image_w=np.size(im, 1), info=info)
def extract_feat(split_idx, img_list, cfg, args, actor: ActorHandle): num_images = len(img_list) print('Number of images on split{}: {}.'.format(split_idx, num_images)) model = DefaultTrainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) model.eval() register() print(cfg.MODEL.BUA.EXTRACTOR.MODE) for im_file in (img_list): # if os.path.exists(os.path.join(special_output_dir, im_file.split('.')[0]+'.npz')): # actor.update.remote(1) # continue image_id = im_file.split('.')[0] # xxx.jpg dump_folder = os.path.join(args.output_dir, str(image_id) + ".npz") if os.path.exists( os.path.join(args.output_dir, im_file.split('.')[0] + '.npz')): actor.update.remote(1) continue # else: # start = True # if not start: # actor.update.remote(1) # continue im = cv2.imread(os.path.join(args.image_dir, im_file)) illegal = False if im is None: illegal = True elif im.shape[-1] != 3: illegal = True elif max(im.shape[:2]) / min(im.shape[:2]) > 10 or max( im.shape[:2]) < 25: illegal = True if illegal: # print(os.path.join(args.image_dir, im_file), "is illegal!") actor.update.remote(1) continue # dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN) pixel_mean = cfg.MODEL.PIXEL_MEAN if args.mode == "caffe" else 0.0 image_h = np.size(im, 0) image_w = np.size(im, 1) dataset_dict = get_image_blob(im, pixel_mean) # extract roi features if cfg.MODEL.BUA.EXTRACTOR.MODE == 1: attr_scores = None with torch.set_grad_enabled(False): if cfg.MODEL.BUA.ATTRIBUTE_ON: boxes, scores, features_pooled, attr_scores = model_inference( model, [dataset_dict], 1, dump_folder, image_h, image_w, True) else: boxes, scores, features_pooled = model_inference( model, [dataset_dict], 1, dump_folder, image_h, image_w) elif cfg.MODEL.BUA.EXTRACTOR.MODE == 2: with torch.set_grad_enabled(False): boxes, scores, _ = model_inference(model, [dataset_dict], 2, dump_folder, image_h, image_w) elif cfg.MODEL.BUA.EXTRACTOR.MODE == 3: # extract roi features by bbox npy = False if os.path.exists( os.path.join(args.bbox_dir, im_file.split('.')[0] + '.npy')): npy = True elif not os.path.exists( os.path.join(args.bbox_dir, im_file.split('.')[0] + '.npz')): actor.update.remote(1) continue if npy: try: bbox = torch.from_numpy( np.load(os.path.join(args.bbox_dir, im_file.split('.')[0] + '.npy'), allow_pickle=True).tolist() ['bbox']) * dataset_dict['im_scale'] except Exception as e: print(e) continue actor.update.remote(1) else: bbox = torch.from_numpy( np.load( os.path.join( args.bbox_dir, im_file.split('.')[0] + '.npz'))['bbox']) * dataset_dict['im_scale'] proposals = Instances(dataset_dict['image'].shape[-2:]) proposals.proposal_boxes = BUABoxes(bbox) dataset_dict['proposals'] = proposals attr_scores = None with torch.set_grad_enabled(False): if cfg.MODEL.BUA.ATTRIBUTE_ON: boxes, scores, features_pooled, attr_scores = model_inference( model, [dataset_dict], 3, dump_folder, image_h, image_w, True) else: boxes, scores, features_pooled = model_inference( model, [dataset_dict], 3, dump_folder, image_h, image_w) else: raise Exception("extract mode not supported.") actor.update.remote(1)
# In[4]: config_file = '../configs/bua-caffe/extract-bua-caffe-r101.yaml' cfg = get_cfg() add_bottom_up_attention_config(cfg, True) cfg.merge_from_file(config_file) cfg.freeze() # In[5]: MIN_BOXES = 10 MAX_BOXES = 20 CONF_THRESH = 0.4 model = DefaultTrainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( '../' + cfg.MODEL.WEIGHTS, resume=True) model.eval() # In[6]: from glob import glob # im_files = ['../datasets/demo/004545.jpg'] im_files = sorted(glob("../../sample_data/*.png")) for im_file in im_files: im = cv2.imread(im_file) dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN)
def main(): parser = argparse.ArgumentParser( description="PyTorch Object Detection2 Inference") parser.add_argument( "--config-file", default="configs/bua-caffe/extract-bua-caffe-r101.yaml", metavar="FILE", help="path to config file", ) parser.add_argument("--mode", default="caffe", type=str, help="bua_caffe, ...") parser.add_argument('--out-dir', dest='output_dir', help='output directory for features', default="data/features") parser.add_argument('--image-dir', dest='image_dir', help='directory with images', default="data/coco_img") parser.add_argument('--gt-bbox-dir', dest='gt_bbox_dir', help='directory with gt-bbox', default="data/bbox") parser.add_argument( "--resume", action="store_true", help="whether to attempt to resume from the checkpoint directory", ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() cfg = setup(args) MIN_BOXES = cfg.MODEL.BUA.EXTRACTOR.MIN_BOXES MAX_BOXES = cfg.MODEL.BUA.EXTRACTOR.MAX_BOXES CONF_THRESH = cfg.MODEL.BUA.EXTRACTOR.CONF_THRESH model = DefaultTrainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) # Extract features. imglist = os.listdir(args.image_dir) num_images = len(imglist) print('Number of images: {}.'.format(num_images)) model.eval() for im_file in tqdm.tqdm(imglist): if os.path.exists( os.path.join(args.output_dir, im_file.split('.')[0] + '.npz')): continue im = cv2.imread(os.path.join(args.image_dir, im_file)) if im is None: print(os.path.join(args.image_dir, im_file), "is illegal!") continue dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN) # extract roi features if cfg.MODEL.BUA.EXTRACTOR.MODE == 1: attr_scores = None with torch.set_grad_enabled(False): if cfg.MODEL.BUA.ATTRIBUTE_ON: boxes, scores, features_pooled, attr_scores = model( [dataset_dict]) else: boxes, scores, features_pooled = model([dataset_dict]) save_roi_features(args, cfg, im_file, im, dataset_dict, boxes, scores, features_pooled, attr_scores) # extract bbox only elif cfg.MODEL.BUA.EXTRACTOR.MODE == 2: with torch.set_grad_enabled(False): boxes, scores = model([dataset_dict]) save_bbox(args, cfg, im_file, im, dataset_dict, boxes, scores) # extract roi features by gt bbox elif cfg.MODEL.BUA.EXTRACTOR.MODE == 3: if not os.path.exists( os.path.join(args.gt_bbox_dir, im_file.split('.')[0] + '.npz')): continue bbox = torch.from_numpy( np.load( os.path.join(args.gt_bbox_dir, im_file.split('.')[0] + '.npz'))['bbox']) * dataset_dict['im_scale'] proposals = Instances(dataset_dict['image'].shape[-2:]) proposals.proposal_boxes = BUABoxes(bbox) dataset_dict['proposals'] = proposals attr_scores = None with torch.set_grad_enabled(False): if cfg.MODEL.BUA.ATTRIBUTE_ON: boxes, scores, features_pooled, attr_scores = model( [dataset_dict]) else: boxes, scores, features_pooled = model([dataset_dict]) save_roi_features_by_gt_bbox(args, cfg, im_file, im, dataset_dict, boxes, scores, features_pooled, attr_scores)
def main(): parser = argparse.ArgumentParser( description="PyTorch Object Detection2 Inference") parser.add_argument( "--config-file", default="configs/bua-caffe/extract-bua-caffe-r101-fix36.yaml", metavar="FILE", help="path to config file", ) # --image-dir or --image parser.add_argument('--image-dir', dest='image_dir', help='directory with images', default="datasets/demos") parser.add_argument( '--image', dest='image', help='image') # e.g. datasets/demos/COCO_val2014_000000060623.jpg parser.add_argument("--mode", default="caffe", type=str, help="bua_caffe, ...") parser.add_argument('--out-dir', dest='output_dir', help='output directory for features', default="features") parser.add_argument('--out-name', dest='output_name', help='output file name for features', default="demos") parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() cfg = setup(args) model = DefaultTrainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=True) model.eval() # Extract features. if args.image: imglist = [args.image] else: imglist = os.listdir(args.image_dir) imglist = [os.path.join(args.image_dir, fn) for fn in imglist] num_images = len(imglist) print('Number of images: {}.'.format(num_images)) imglist.sort() MIN_BOXES = cfg.MODEL.BUA.EXTRACTOR.MIN_BOXES MAX_BOXES = cfg.MODEL.BUA.EXTRACTOR.MAX_BOXES CONF_THRESH = cfg.MODEL.BUA.EXTRACTOR.CONF_THRESH classes = [] with open(os.path.join('evaluation/objects_vocab.txt')) as f: for object in f.readlines(): names = [n.lower().strip() for n in object.split(',')] classes.append(names[0]) attributes = [] with open(os.path.join('evaluation/attributes_vocab.txt')) as f: for att in f.readlines(): names = [n.lower().strip() for n in att.split(',')] attributes.append(names[0]) classes = np.array(classes) attributes = np.array(attributes) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) with h5py.File(os.path.join(args.output_dir, '%s_fc.h5' % args.output_name), 'a') as file_fc, \ h5py.File(os.path.join(args.output_dir, '%s_att.h5' % args.output_name), 'a') as file_att, \ h5py.File(os.path.join(args.output_dir, '%s_box.h5' % args.output_name), 'a') as file_box: informations = {} try: for im_file in tqdm.tqdm(imglist): img_nm = os.path.basename(im_file) im = cv2.imread(im_file) if im is None: print(im_file, "is illegal!") continue dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN) # extract roi features attr_scores = None with torch.set_grad_enabled(False): if cfg.MODEL.BUA.ATTRIBUTE_ON: boxes, scores, features_pooled, attr_scores = model( [dataset_dict]) else: boxes, scores, features_pooled = model([dataset_dict]) dets = boxes[0].tensor.cpu() / dataset_dict['im_scale'] scores = scores[0].cpu() feats = features_pooled[0].cpu() max_conf = torch.zeros((scores.shape[0])).to(scores.device) for cls_ind in range(1, scores.shape[1]): cls_scores = scores[:, cls_ind] keep = nms(dets, cls_scores, 0.3) max_conf[keep] = torch.where( cls_scores[keep] > max_conf[keep], cls_scores[keep], max_conf[keep]) keep_boxes = torch.nonzero(max_conf >= CONF_THRESH).flatten() if len(keep_boxes) < MIN_BOXES: keep_boxes = torch.argsort(max_conf, descending=True)[:MIN_BOXES] elif len(keep_boxes) > MAX_BOXES: keep_boxes = torch.argsort(max_conf, descending=True)[:MAX_BOXES] image_feat = feats[keep_boxes].numpy() image_bboxes = dets[keep_boxes].numpy() image_objects_conf = np.max(scores[keep_boxes].numpy()[:, 1:], axis=1) image_objects = classes[np.argmax( scores[keep_boxes].numpy()[:, 1:], axis=1)] info = { 'image_name': img_nm, 'image_h': np.size(im, 0), 'image_w': np.size(im, 1), 'num_boxes': len(keep_boxes), 'objects': image_objects, 'objects_conf': image_objects_conf } if attr_scores is not None: attr_scores = attr_scores[0].cpu() image_attrs_conf = np.max( attr_scores[keep_boxes].numpy()[:, 1:], axis=1) image_attrs = attributes[np.argmax( attr_scores[keep_boxes].numpy()[:, 1:], axis=1)] info['attrs'] = image_attrs info['attrs_conf'] = image_attrs_conf file_fc.create_dataset(img_nm, data=image_feat.mean(0)) file_att.create_dataset(img_nm, data=image_feat) file_box.create_dataset(img_nm, data=image_bboxes) informations[img_nm] = info finally: file_fc.close() file_att.close() file_box.close() pickle.dump( informations, open( os.path.join(args.output_dir, '%s_info.pkl' % args.output_name), 'wb')) print( '--------------------------------------------------------------------' )
def test_equivalence(cfg, infile, outfile): from deeplab3.modeling.backbone.resnet import ResNet101 from detectron2.checkpoint.detection_checkpoint import DetectionCheckpointer from detectron2.engine import DefaultTrainer import torch.nn as nn #Deeplab model = ResNet101(cfg.MODEL.DEEPLAB, BatchNorm=nn.BatchNorm2d) model.eval() if torch.cuda.is_available(): checkpoint = torch.load(infile, map_location=torch.device('cuda')) else: checkpoint = torch.load(infile, map_location=torch.device('cpu')) state_dict = { k.replace('backbone.', ''): v for k, v in checkpoint['state_dict'].items() if 'backbone' in k } model.load_state_dict(state_dict) #Detectron2 model2 = DefaultTrainer.build_model(cfg) model2.eval() DetectionCheckpointer(model2, save_dir=outfile).resume_or_load(cfg.MODEL.WEIGHTS) #Check that the layers are identical convs = [] downsample = [] bn = [] downsample_bn = [] for name, param in model.state_dict().items(): param_dict = {'name': name, 'param': param, 'size': param.size()} if 'conv' in name: convs.append(param_dict) elif 'downsample.0' in name: downsample.append(param_dict) elif 'bn' in name: bn.append(param_dict) elif 'downsample.1' in name: downsample_bn.append(param_dict) else: print(name) convs2 = [] shortcut = [] bn2 = [] downsample_bn2 = [] for name, param in model2.backbone.state_dict().items(): param_dict = {'name': name, 'param': param, 'size': param.size()} if 'conv' in name: convs.append(param_dict) elif 'downsample.0' in name: downsample.append(param_dict) elif 'bn' in name: bn.append(param_dict) elif 'downsample.1' in name: downsample_bn.append(param_dict) else: print(name) for p1, p2 in zip(convs, convs2): print(p1['name'], p2['name']) if (p1['size'] != p2['size']): print(p1['size'], p2['size']) raise ValueError('Sizes not equal') assert (torch.all(torch.eq(p1['param'], p2['param']))) for p1, p2 in zip(shortcut, downsample): print(p1['name'], p2['name']) if (p1['size'] != p2['size']): print(p1['size'], p2['size']) raise ValueError('Sizes not equal') assert (torch.all(torch.eq(p1['param'], p2['param']))) for p1, p2 in zip(bn, bn2): print(p1['name'], p2['name']) if (p1['size'] != p2['size']): print(p1['size'], p2['size']) raise ValueError('Sizes not equal') assert (torch.all(torch.eq(p1['param'], p2['param']))) for p1, p2 in zip(downsample_bn2, downsample_bn): print(p1['name'], p2['name']) if (p1['size'] != p2['size']): print(p1['size'], p2['size']) raise ValueError('Sizes not equal') assert (torch.all(torch.eq(p1['param'], p2['param']))) #Check that output is identical input = torch.ones((1, 4, 640, 640), dtype=torch.float32) output = model(input) output2 = model2.backbone(input) assert (torch.all(torch.eq(output['res5'], output2['res5']))) assert (torch.all(torch.eq(output['res2'], output2['res2']))) print("Passed all tests")