def save_bbox(args, cfg, im_file, im, dataset_dict, boxes, scores):
MIN_BOXES = cfg.MODEL.BUA.EXTRACTOR.MIN_BOXES
MAX_BOXES = cfg.MODEL.BUA.EXTRACTOR.MAX_BOXES
CONF_THRESH = cfg.MODEL.BUA.EXTRACTOR.CONF_THRESH
scores = scores[0].cpu()
boxes = boxes[0]
num_classes = scores.shape[1]
boxes = BUABoxes(boxes.reshape(-1, 4))
boxes.clip((dataset_dict['image'].shape[1]/dataset_dict['im_scale'], dataset_dict['image'].shape[2]/dataset_dict['im_scale']))
boxes = boxes.tensor.view(-1, num_classes*4).cpu() # R x C x 4
cls_boxes = torch.zeros((boxes.shape[0], 4))
for idx in range(boxes.shape[0]):
cls_idx = torch.argmax(scores[idx, 1:]) + 1
cls_boxes[idx, :] = boxes[idx, cls_idx * 4:(cls_idx + 1) * 4]
max_conf = torch.zeros((scores.shape[0])).to(scores.device)
for cls_ind in range(1, num_classes):
cls_scores = scores[:, cls_ind]
keep = nms(cls_boxes, cls_scores, 0.3)
max_conf[keep] = torch.where(cls_scores[keep] > max_conf[keep],
cls_scores[keep],
max_conf[keep])
keep_boxes = torch.argsort(max_conf, descending=True)[:MAX_BOXES]
image_bboxes = cls_boxes[keep_boxes]
output_file = os.path.join(args.output_dir, im_file.split('.')[0])
np.savez_compressed(output_file, bbox=image_bboxes, num_bbox=len(keep_boxes), image_h=np.size(im, 0), image_w=np.size(im, 1))
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="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)