def main(name_scope, gpu_dev, num_images, args):
t=args.t
model = initialize_model_from_cfg()
num_classes = cfg.MODEL.NUM_CLASSES
all_boxes, all_segms, all_keyps = empty_results(num_classes, num_images)
temp_frame_folder = osp.join(args.out_path,args.vid_name + '_frames/',str(t))
imgs = glob.glob(temp_frame_folder+'/*.jpg')
for i in range(len(imgs)):
if i%100==0:
print('Processing Detection for Frame %d'%(i+1))
im_ = cv2.imread(imgs[i])
assert im_ is not None
im_ = np.expand_dims(im_, 0)
with core.NameScope(name_scope):
with core.DeviceScope(gpu_dev):
cls_boxes_i, cls_segms_i, cls_keyps_i = im_detect_all(
model, im_, None) #TODO: Parallelize detection
extend_results(i, all_boxes, cls_boxes_i)
if cls_segms_i is not None:
extend_results(i, all_segms, cls_segms_i)
if cls_keyps_i is not None:
extend_results(i, all_keyps, cls_keyps_i)
det_name = args.vid_name + '_' + str(args.t) + '_detections.pkl'
det_file = osp.join(args.out_path, det_name)
robust_pickle_dump(dict(all_keyps=all_keyps),det_file)
shutil.rmtree(osp.join(args.out_path, args.vid_name + '_frames'))
def im_list_detections(model, roidb):
"""Generate RetinaNet detections on all images in an imdb."""
_t = Timer()
num_images = len(roidb)
num_classes = cfg.MODEL.NUM_CLASSES
all_boxes, _, _ = test_engine.empty_results(num_classes, num_images)
# create anchors for each level
anchors = create_cell_anchors()
for i, entry in enumerate(roidb):
im = cv2.imread(entry['image'])
with c2_utils.NamedCudaScope(0):
_t.tic()
cls_boxes_i = im_detections(model, im, anchors)
_t.toc()
test_engine.extend_results(i, all_boxes, cls_boxes_i)
logger.info('im_detections: {:d}/{:d} {:.3f}s'.format(
i + 1, num_images, _t.average_time))
return all_boxes
else: # For subprocess call
assert cfg.TEST.DATASETS, 'cfg.TEST.DATASETS shouldn\'t be empty'
assert_and_infer_cfg()
logger.info('Re-evaluation with config:')
logger.info(pprint.pformat(cfg))
with open(args.result_path, 'rb') as f:
results = pickle.load(f)
logger.info('Loading results from {}.'.format(args.result_path))
all_boxes = results['all_boxes']
dataset_name = cfg.TEST.DATASETS[0]
dataset = JsonDataset(dataset_name)
roidb = dataset.get_roidb()
num_images = len(roidb)
num_classes = cfg.MODEL.NUM_CLASSES + 1
final_boxes = empty_results(num_classes, num_images)
test_corloc = 'train' in dataset_name
for i, entry in enumerate(roidb):
boxes = all_boxes[entry['image']]
if test_corloc:
_, _, cls_boxes_i = box_results_for_corloc(boxes['scores'],
boxes['boxes'])
else:
_, _, cls_boxes_i = box_results_with_nms_and_limit(
boxes['scores'], boxes['boxes'])
extend_results(i, final_boxes, cls_boxes_i)
results = task_evaluation.evaluate_all(dataset, final_boxes,
args.output_dir, test_corloc)
def main():
"""Main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
logger = utils.logging.setup_logging(__name__)
args = parse_args()
logger.info('Called with args:')
logger.info(args)
assert (torch.cuda.device_count() == 1) ^ bool(args.multi_gpu_testing)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
if args.output_dir is None:
ckpt_path = args.load_ckpt if args.load_ckpt else args.load_detectron
args.output_dir = os.path.join(
os.path.dirname(os.path.dirname(ckpt_path)), 'test')
logger.info('Automatically set output directory to %s', args.output_dir)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
if args.close_fpn:
args.cfg_file = "configs/few_shot/e2e_mask_rcnn_R-50-C4_1x_{}.yaml".format(args.group)
else:
args.cfg_file = "configs/few_shot/e2e_mask_rcnn_R-50-FPN_1x_{}.yaml".format(args.group)
if args.cfg_file is not None:
merge_cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
merge_cfg_from_list(args.set_cfgs)
cfg.VIS = args.vis
cfg.SEEN = args.seen
if args.close_co_atten:
cfg.CO_ATTEN = False
if args.close_relation_rcnn:
cfg.RELATION_RCNN = False
if not args.close_fpn:
cfg.FAST_RCNN.ROI_BOX_HEAD = 'fast_rcnn_heads.roi_2mlp_head'
cfg.MRCNN.ROI_MASK_HEAD = 'mask_rcnn_heads.mask_rcnn_fcn_head_v1up4convs'
else:
cfg.FAST_RCNN.ROI_BOX_HEAD = 'torchResNet.ResNet_roi_conv5_head'
cfg.MRCNN.ROI_MASK_HEAD = 'mask_rcnn_heads.mask_rcnn_fcn_head_v0upshare'
if args.deform_conv:
cfg.MODEL.USE_DEFORM = True
if args.dataset == "fis_cell":
cfg.TEST.DATASETS = ('fis_cell_test',)
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "coco2017":
cfg.TEST.DATASETS = ('coco_2017_val',)
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == "keypoints_coco2017":
cfg.TEST.DATASETS = ('keypoints_coco_2017_val',)
cfg.MODEL.NUM_CLASSES = 2
else: # For subprocess call
assert cfg.TEST.DATASETS, 'cfg.TEST.DATASETS shouldn\'t be empty'
assert_and_infer_cfg()
#logger.info('Testing with config:')
#logger.info(pprint.pformat(cfg))
# manually set args.cuda
args.cuda = True
timer_for_ds = defaultdict(Timer)
### Dataset ###
timer_for_ds['roidb'].tic()
imdb, roidb, ratio_list, ratio_index, query, cat_list = combined_roidb(
cfg.TEST.DATASETS, False)
timer_for_ds['roidb'].toc()
roidb_size = len(roidb)
logger.info('{:d} roidb entries'.format(roidb_size))
logger.info('Takes %.2f sec(s) to construct roidb', timer_for_ds['roidb'].average_time)
batchSampler = BatchSampler(
sampler=MinibatchSampler(ratio_list, ratio_index, shuffle=False),
batch_size=1,
drop_last=False
)
dataset = RoiDataLoader(
roidb, ratio_list, ratio_index, query,
cfg.MODEL.NUM_CLASSES,
training=False, cat_list=cat_list, shot=args.checkshot)
### Model ###
model = initialize_model_from_cfg(args, gpu_id=0)
all_results = OrderedDict({
'box':
OrderedDict(
[
('AP', []),
('AP50', []),
('AP75', []),
('APs', []),
('APm', []),
('APl', []),
]
),
'mask':
OrderedDict(
[
('AP', []),
('AP50', []),
('AP75', []),
('APs', []),
('APm', []),
('APl', []),
]
)
})
timer_for_total = defaultdict(Timer)
timer_for_total['total_test_time'].tic()
for avg in range(args.average):
dataset.query_position = avg
dataloader = torch.utils.data.DataLoader(
dataset,
batch_sampler=batchSampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch
)
dataiterator = iter(dataloader)
num_images = len(ratio_index)
num_cats = imdb.num_classes
all_boxes, all_segms, all_keyps = empty_results(num_cats, num_images)
# total quantity of testing images
num_detect = len(ratio_index)
timers = defaultdict(Timer)
post_fix = '%dshot_g%d_seen%d_%d'%(args.checkshot, args.group, args.seen, avg)
for i, index in enumerate(ratio_index):
input_data = next(dataiterator)
catgory = input_data['choice']
entry = dataset._roidb[dataset.ratio_index[i]]
if cfg.TEST.PRECOMPUTED_PROPOSALS:
# The roidb may contain ground-truth rois (for example, if the roidb
# comes from the training or val split). We only want to evaluate
# detection on the *non*-ground-truth rois. We select only the rois
# that have the gt_classes field set to 0, which means there's no
# ground truth.
box_proposals = entry['boxes'][entry['gt_classes'] == 0]
if len(box_proposals) == 0:
continue
else:
# Faster R-CNN type models generate proposals on-the-fly with an
# in-network RPN; 1-stage models don't require proposals.
box_proposals = None
#im = cv2.imread(entry['image'])
im = imread(entry['image'])
if len(im.shape) == 2:
im = im[:,:,np.newaxis]
im = np.concatenate((im,im,im), axis=2)
alpha = 2.2 # Simple contrast control
beta = 0.5 # Simple brightness control
im_colored = im.copy()
if catgory[0].item() == 7:
im_colored[:,:,0] = 0
im_colored[:,:,1] = 0
im_colored = cv2.convertScaleAbs(im_colored, alpha=alpha, beta=beta)
elif catgory[0].item() == 8:
im_colored[:,:,1] = 0
im_colored[:,:,2] = 0
im_colored = cv2.convertScaleAbs(im_colored, alpha=alpha, beta=beta)
elif catgory[0].item() == 6:
im_colored[:,:,0] = 0
im_colored[:,:,2] = 0
im_colored = cv2.convertScaleAbs(im_colored, alpha=alpha, beta=beta)
cls_boxes_i, cls_segms_i, cls_keyps_i = im_detect_all(model, im, input_data['query'], input_data['query_type'], catgory, num_cats, box_proposals, timers)
im = im_colored
extend_results(i, all_boxes, cls_boxes_i)
if cls_segms_i is not None:
extend_results(i, all_segms, cls_segms_i)
if cls_keyps_i is not None:
extend_results(i, all_keyps, cls_keyps_i)
if i % 10 == 0: # Reduce log file size
ave_total_time = np.sum([t.average_time for t in timers.values()])
eta_seconds = ave_total_time * (num_images - i - 1)
eta = str(datetime.timedelta(seconds=int(eta_seconds)))
det_time = (
timers['im_detect_bbox'].average_time +
timers['im_detect_mask'].average_time +
timers['im_detect_keypoints'].average_time
)
misc_time = (
timers['misc_bbox'].average_time +
timers['misc_mask'].average_time +
timers['misc_keypoints'].average_time
)
logger.info(
(
'im_detect: range [{:d}, {:d}] of {:d}: '
'{:d}/{:d} {:.3f}s + {:.3f}s (eta: {})'
).format(
1, num_detect, num_detect, i + 1,
num_detect, det_time, misc_time, eta
)
)
if cfg.VIS:
im_name = entry['image']
class_name = im_name.split('/')[-4]
file_name = im_name.split('/')[-3]
im_target = im.copy()
to_tensor = transforms.ToTensor()
o_querys=[]
for i in range(args.checkshot):
o_query = input_data['query'][0][i][0].permute(1, 2,0).contiguous().cpu().numpy()
o_query *= [0.229, 0.224, 0.225]
o_query += [0.485, 0.456, 0.406]
o_query *= 255
o_query_colored = o_query.copy()
if catgory[0].item() == 7:
o_query_colored[:,:,0] = 0
o_query_colored[:,:,1] = 0
o_query_colored = cv2.convertScaleAbs(o_query_colored, alpha=alpha, beta=beta)
elif catgory[0].item() == 8:
o_query_colored[:,:,1] = 0
o_query_colored[:,:,2] = 0
o_query_colored = cv2.convertScaleAbs(o_query_colored, alpha=alpha, beta=beta)
elif catgory[0].item() == 6:
o_query_colored[:,:,0] = 0
o_query_colored[:,:,2] = 0
o_query_colored = cv2.convertScaleAbs(o_query_colored, alpha=alpha, beta=beta)
o_query = o_query_colored
o_query = Image.fromarray(o_query.astype(np.uint8))
o_querys.append(to_tensor(o_query))
o_querys_grid = make_grid(o_querys, nrow=args.checkshot//2, normalize=True, scale_each=True, pad_value=1)
o_querys_grid = transforms.ToPILImage()(o_querys_grid).convert("RGB")
query_w, query_h = o_querys_grid.size
query_bg = Image.new('RGB', (im_target.shape[1], im_target.shape[0]), (255, 255, 255))
bg_w, bg_h = query_bg.size
offset = ((bg_w - query_w) // 2, (bg_h - query_h) // 2)
query_bg.paste(o_querys_grid, offset)
query = np.asarray(query_bg)
im_pair = np.concatenate((im_target, query), axis=1)
im_output_dir = os.path.join(args.output_dir, 'vis', post_fix, class_name)
if not os.path.exists(im_output_dir):
os.makedirs(im_output_dir)
sample_output_dir = os.path.join(im_output_dir, os.path.basename('{:d}_{:s}'.format(i, file_name)))
if not os.path.exists(sample_output_dir):
os.makedirs(sample_output_dir)
target_save_name = os.path.join(sample_output_dir, os.path.basename('{:d}_{:s}'.format(i, file_name)) + '_target.pdf')
target = Image.fromarray(im_target.astype(np.uint8))
target.save(target_save_name,"pdf")
query_save_name = os.path.join(sample_output_dir, os.path.basename('{:d}_{:s}'.format(i, file_name)) + '_query.pdf')
query = Image.fromarray(query.astype(np.uint8))
query.save(query_save_name,"pdf")
pred_save_name = os.path.join(sample_output_dir, os.path.basename('{:d}_{:s}'.format(i, file_name)) + '_pred.pdf')
vis_utils.save_one_image(
im,
pred_save_name,
cls_boxes_i,
segms = cls_segms_i,
keypoints = cls_keyps_i,
thresh = cfg.VIS_TH,
box_alpha = 0.6,
dataset = imdb,
show_class = False
)
im_det = vis_utils.vis_one_image(
im,
'{:d}_det_{:s}'.format(i, file_name),
os.path.join(args.output_dir, 'vis', post_fix),
cls_boxes_i,
segms = cls_segms_i,
keypoints = cls_keyps_i,
thresh = cfg.VIS_TH,
box_alpha = 0.6,
dataset = imdb,
show_class = False,
class_name = class_name,
draw_bbox = False
)
gt_save_name = os.path.join(sample_output_dir, os.path.basename('{:d}_{:s}'.format(i, file_name)) + '_gt.pdf')
vis_utils.save_one_image_gt(
im, entry['id'],
gt_save_name,
dataset = imdb)
im_gt = vis_utils.vis_one_image_gt(
im, entry['id'],
'{:d}_gt_{:s}'.format(i, file_name),
os.path.join(args.output_dir, 'vis', post_fix),
dataset = imdb,
class_name = class_name)
im_det = np.asarray(im_det)
im_gt = np.asarray(im_gt)
im2draw = np.concatenate((im_gt, im_det), axis=1)
im2show = np.concatenate((im_pair, im2draw), axis=0)
im_save_name = os.path.basename('{:d}_{:s}'.format(i, file_name)) + '.png'
cv2.imwrite(os.path.join(im_output_dir, '{}'.format(im_save_name)), cv2.cvtColor(im2show, cv2.COLOR_RGB2BGR))
"""
if cfg.VIS:
im_name = dataset._roidb[dataset.ratio_index[i]]['image']
class_name = im_name.split('/')[-4]
file_name = im_name.split('/')[-3]
im2show = Image.open(im_name).convert("RGB")
o_query = input_data['query'][0][0][0].permute(1, 2, 0).contiguous().cpu().numpy()
o_query *= [0.229, 0.224, 0.225]
o_query += [0.485, 0.456, 0.406]
o_query *= 255
o_query = o_query[:,:,::-1]
o_query = Image.fromarray(o_query.astype(np.uint8))
query_w, query_h = o_query.size
query_bg = Image.new('RGB', (im2show.size), (255, 255, 255))
bg_w, bg_h = query_bg.size
offset = ((bg_w - query_w) // 2, (bg_h - query_h) // 2)
query_bg.paste(o_query, offset)
im2show = np.asarray(im2show)
o_query = np.asarray(query_bg)
im2show = np.concatenate((im2show, o_query), axis=1)
output_dir = os.path.join(args.output_dir, 'vis')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
im_save_dir = os.path.join(output_dir, class_name)
if not os.path.exists(im_save_dir):
os.makedirs(im_save_dir)
im_save_name = os.path.join(im_save_dir, file_name + '_%d_d.png'%(i))
cv2.imwrite(im_save_name, cv2.cvtColor(im2show, cv2.COLOR_RGB2BGR))
"""
cfg_yaml = yaml.dump(cfg)
#det_name = 'detections.pkl'
det_file = os.path.join(args.output_dir, 'detections_' + post_fix + '.pkl')
save_object(
dict(
all_boxes=all_boxes,
all_segms=all_segms,
all_keyps=all_keyps,
cfg=cfg_yaml
), det_file
)
logger.info('Wrote detections to: {}'.format(os.path.abspath(det_file)))
results = task_evaluation.evaluate_all(
imdb, all_boxes, all_segms, all_keyps, args.output_dir
)
task_evaluation.check_expected_results(
results,
atol=cfg.EXPECTED_RESULTS_ATOL,
rtol=cfg.EXPECTED_RESULTS_RTOL
)
for task, metrics in all_results.items():
metric_names = metrics.keys()
for metric_name in metric_names:
all_results[task][metric_name].append(results[imdb.name][task][metric_name])
#task_evaluation.log_copy_paste_friendly_results(results)
for task, metrics in all_results.items():
metric_names = metrics.keys()
for metric_name in metric_names:
values = all_results[task][metric_name]
all_results[task][metric_name] = sum(values) / len(values)
post_fix = '%dshot_g%d_seen%d'%(args.checkshot, args.group, args.seen)
avg_results_path = os.path.join(args.output_dir, ('avg_cocoeval_' + post_fix + '_results.json'))
with open(avg_results_path, 'w') as f:
f.write(json.dumps(all_results))
timer_for_total['total_test_time'].toc()
logger.info('Total inference time: {:.3f}s'.format(timer_for_total['total_test_time'].average_time))
def main(name_scope, gpu_dev, num_images, args):
model = initialize_model_from_cfg()
num_classes = cfg.MODEL.NUM_CLASSES
all_boxes, all_segms, all_keyps = empty_results(num_classes, num_images)
if '2d_best' in args.cfg_file:
for i in range(num_images):
print('Processing Detection for Frame %d'%(i+1))
im_ = _read_video_frames(args.out_path, args.vid_name, i)
im_ = np.expand_dims(im_, 0)
with core.NameScope(name_scope):
with core.DeviceScope(gpu_dev):
cls_boxes_i, cls_segms_i, cls_keyps_i = im_detect_all(
model, im_, None) #TODO: Parallelize detection
extend_results(i, all_boxes, cls_boxes_i)
if cls_segms_i is not None:
extend_results(i, all_segms, cls_segms_i)
if cls_keyps_i is not None:
extend_results(i, all_keyps, cls_keyps_i)
elif '3d' in args.cfg_file:
for i in range(num_images-2):
print('Processing Detection for Frame %d to Frame %d' % (i + 1, i + 2))
ims_ = _read_video_3frames(args.out_path, args.vid_name, i)
# ims_ = np.expand_dims(ims_, 0)
with core.NameScope(name_scope):
with core.DeviceScope(gpu_dev):
cls_boxes_i, cls_segms_i, cls_keyps_i = im_detect_all(
model, ims_, None)
# extend boxes for 3 frames
tmp_boxes_i2 = deepcopy(cls_boxes_i)
tmp_boxes_i2[1] = tmp_boxes_i2[1][:, 8:]
extend_results(i+2, all_boxes, tmp_boxes_i2)
tmp_boxes_i1 = deepcopy(cls_boxes_i)
tmp_boxes_i1[1] = tmp_boxes_i1[1][:, [4, 5, 6, 7, -1]]
extend_results(i+1, all_boxes, tmp_boxes_i1)
tmp_boxes_i0 = deepcopy(cls_boxes_i)
tmp_boxes_i0[1] = tmp_boxes_i0[1][:, [0, 1, 2, 3, -1]]
extend_results(i, all_boxes, tmp_boxes_i0)
# extend segms for 3 frames
if cls_segms_i is not None:
extend_results(i+2, all_segms, cls_segms_i)
# extend keyps for 3 frames
if cls_keyps_i is not None:
# extend the i+2 th one
tmp_keyps_i2 = deepcopy(cls_keyps_i)
for idx in range(len(tmp_keyps_i2[1])):
tmp_keyps_i2[1][idx] = tmp_keyps_i2[1][idx][:, 34:]
extend_results(i+2, all_keyps, tmp_keyps_i2)
# extend the i+1 th one
tmp_keyps_i1 = deepcopy(cls_keyps_i)
for idx in range(len(tmp_keyps_i1[1])):
tmp_keyps_i1[1][idx] = tmp_keyps_i1[1][idx][:, 17:34]
extend_results(i + 1, all_keyps, tmp_keyps_i1)
# extend the i th one
tmp_keyps_i0 = deepcopy(cls_keyps_i)
for idx in range(len(tmp_keyps_i0[1])):
tmp_keyps_i0[1][idx] = tmp_keyps_i0[1][idx][:, :17]
extend_results(i, all_keyps, tmp_keyps_i0)
cfg_yaml = yaml.dump(cfg)
det_name = args.vid_name + '_detections.pkl'
det_file = osp.join(args.out_path, det_name)
robust_pickle_dump(
dict(all_boxes=all_boxes,
all_segms=all_segms,
all_keyps=all_keyps,
cfg=cfg_yaml),
det_file)
frames = sorted(glob.glob(osp.join(args.out_path,args.vid_name + '_frames','*.jpg')))
out_detrack_file = osp.join(args.out_path, args.vid_name + '_detections_withTracks.pkl')
# Debug configurations
if cfg.TRACKING.DEBUG.UPPER_BOUND_2_GT_KPS: # if this is true
cfg.TRACKING.DEBUG.UPPER_BOUND = True # This must be set true
# Set include_gt True when using the roidb to evalute directly. Not doing
# that currently
dets = _load_det_file(det_file)
if cfg.TRACKING.KEEP_CENTER_DETS_ONLY:
_center_detections(dets)
conf = cfg.TRACKING.CONF_FILTER_INITIAL_DETS
dets = _prune_bad_detections(dets, conf)
if cfg.TRACKING.LSTM_TEST.LSTM_TRACKING_ON:
# Needs torch, only importing if we need to run LSTM tracking
from lstm.lstm_track import lstm_track_utils
lstm_model = lstm_track_utils.init_lstm_model(
cfg.TRACKING.LSTM_TEST.LSTM_WEIGHTS)
lstm_model.cuda()
else:
lstm_model = None
dets_withTracks = compute_matches_tracks(frames, dets, lstm_model)
_write_det_file(dets_withTracks, out_detrack_file)
for i in range(num_images):
vis_im = _generate_visualizations(frames[i], i, dets['all_boxes'], dets['all_keyps'], dets['all_tracks'])
cv2.imwrite(osp.join(args.out_path, args.vid_name + '_vis','%08d.jpg'%(i+1)),vis_im)