def generate_npz(extract_mode, *args):
if extract_mode == 1:
save_roi_features(*args)
elif extract_mode == 2:
save_bbox(*args)
elif extract_mode == 3:
save_roi_features_by_bbox(*args)
else:
print('Invalid Extract Mode! ')
def generate_npz(extract_mode, pba: ActorHandle, *args):
if extract_mode == 1:
save_roi_features(*args)
elif extract_mode == 2:
save_bbox(*args)
elif extract_mode == 3:
save_roi_features_by_bbox(*args)
else:
print('Invalid Extract Mode! ')
pba.update.remote(1)
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)