def transform(im_dir, save_path):
trainval_im_names, trainval_labels, _ = transform_trainval_set(im_dir)
test_im_names, test_ids, test_cams, test_marks = transform_test_set(im_dir)
partitions = {
'trainval_im_names': trainval_im_names,
'trainval_labels': trainval_labels,
'test_im_names': test_im_names,
'test_ids': test_ids,
'test_cams': test_cams,
'test_marks': test_marks
}
save_pickle(partitions, save_path)
im_name_to_kpt = {}
im_name_to_h_w = {}
ann_dir = '/mnt/data-1/data/houjing.huang/Dataset/coco/annotations'
im_root_dir = '/mnt/data-1/data/houjing.huang/Dataset/coco'
# 100403 annotations, 39210 with dp_masks, 39210 with keypoints
im_name_to_kpt_, im_name_to_h_w_ = transform(
load_json(osp.join(ann_dir, 'densepose_coco_2014_train.json')),
osp.join(im_root_dir, 'train2014'))
im_name_to_kpt.update(im_name_to_kpt_)
im_name_to_h_w.update(im_name_to_h_w_)
# 25053 annotations, 7297 with dp_masks, 7297 with keypoints
im_name_to_kpt_, im_name_to_h_w_ = transform(
load_json(osp.join(ann_dir, 'densepose_coco_2014_valminusminival.json')),
osp.join(im_root_dir, 'val2014'))
im_name_to_kpt.update(im_name_to_kpt_)
im_name_to_h_w.update(im_name_to_h_w_)
# 5673 annotations, 2243 with dp_masks, 2243 with keypoints
im_name_to_kpt_, im_name_to_h_w_ = transform(
load_json(osp.join(ann_dir, 'densepose_coco_2014_minival.json')),
osp.join(im_root_dir, 'val2014'))
im_name_to_kpt.update(im_name_to_kpt_)
im_name_to_h_w.update(im_name_to_h_w_)
# Done, 48750 im_name to kpt mappings
print('Done, {} im_name to kpt mappings'.format(len(im_name_to_kpt)))
save_pickle(
im_name_to_kpt,
'/mnt/data-1/data/houjing.huang/Dataset/coco_part/im_name_to_kpt.pkl')
save_pickle(
im_name_to_h_w,
'/mnt/data-1/data/houjing.huang/Dataset/coco_part/im_name_to_h_w.pkl')
def main():
cfg = parse_args()
ckpt_file = 'exp/{}_train_Baseline/ckpt.pth'.format(cfg.train_set)
exp_dir = 'exp/{}_sw_occlusion'.format(cfg.train_set)
# Redirect logs to both console and file.
ReDirectSTD(osp.join(exp_dir, 'stdout_{}.txt'.format(time_str())),
'stdout', False)
ReDirectSTD(osp.join(exp_dir, 'stderr_{}.txt'.format(time_str())),
'stderr', False)
TVT, TMO = set_devices(cfg.sys_device_ids)
# Dump the configurations to log.
import pprint
print('-' * 60)
print('cfg.__dict__')
pprint.pprint(cfg.__dict__)
print('-' * 60)
###########
# Dataset #
###########
im_mean = [0.486, 0.459, 0.408]
im_std = [0.229, 0.224, 0.225]
dataset_kwargs = dict(
resize_h_w=cfg.resize_h_w,
scale=True,
im_mean=im_mean,
im_std=im_std,
batch_dims='NCHW',
num_prefetch_threads=cfg.num_prefetch_threads,
prefetch_size=cfg.prefetch_size,
)
# batch_size=1, final_batch=True, mirror_type=None
train_set_kwargs = dict(
name=cfg.train_set,
part='trainval',
batch_size=1,
final_batch=True,
shuffle=True,
crop_prob=cfg.crop_prob,
crop_ratio=cfg.crop_ratio,
mirror_type=None,
prng=np.random,
)
train_set_kwargs.update(dataset_kwargs)
train_set = create_dataset(**train_set_kwargs)
#########
# Model #
#########
model = Model(
last_conv_stride=cfg.last_conv_stride,
max_or_avg=cfg.max_or_avg,
num_classes=len(set(train_set.labels)),
)
# Model wrapper
model_w = DataParallel(model)
ckpt = torch.load(ckpt_file, map_location=(lambda storage, loc: storage))
model.load_state_dict(ckpt['state_dicts'][0])
print('Loaded model weights from {}'.format(ckpt_file))
model.eval()
# Transfer Models to Specified Device
TMO([model])
############################
# Sliding Window Occlusion #
############################
l = int(np.sqrt(cfg.sw_area * np.prod(cfg.resize_h_w)))
sw_h_w = [l, l]
all_masks = gen_masks(cfg.resize_h_w, sw_h_w, cfg.sw_stride)
h_pos, w_pos = get_sw_positions(cfg.resize_h_w, sw_h_w, cfg.sw_stride)
print('Num of all possible masks: {} * {} = {}'.format(
len(h_pos), len(w_pos), len(all_masks)))
def sw_occlude():
"""Calculate the probability difference caused by occluding different positions."""
im_names, prob_diff = [], []
epoch_done = False
num_ims = 0
st_time = time.time()
last_time = time.time()
# For each image
while not epoch_done:
num_ims += 1
im, im_name, label, _, epoch_done = train_set.next_batch()
im_names.append(im_name[0])
im = Variable(TVT(torch.from_numpy(im).float()))
label = label[0]
# Calculate the original prob.
# feat, logits = model_w(im)
# ori_prob = F.softmax(logits, 1).data.cpu().numpy()[0, label]
# To save time, here just use 1.
ori_prob = 1
probs = []
# In order not to over flood the GPU memory, split into small batches.
for masks in np.array_split(all_masks,
int(len(all_masks) / 32) + 1):
# Repeat an image for num_masks times.
# `im` with shape [1, C, H, W] => `repeated_im` with shape [num_masks, C, H, W]
repeated_im = im.repeat(len(masks), 1, 1, 1)
# Repeat each mask for C times.
# `masks` shape [num_masks, H, W] => [num_masks, C, H, W]
masks = Variable(
TVT(
torch.from_numpy(masks).float().unsqueeze(1).expand_as(
repeated_im)))
# `logits` with shape [num_masks, num_classes]
feat, logits = model_w(repeated_im * masks)
probs_ = F.softmax(logits, 1)
probs_ = probs_.data.cpu().numpy()[:, label].flatten()
probs.append(probs_)
# with shape [num_h_pos, num_w_pos], it can be resized to im shape for visualization
probs = np.reshape(np.concatenate(probs), [len(h_pos), len(w_pos)])
prob_diff.append(ori_prob - probs)
if num_ims % 50 == 0:
print('\t{}/{} images done, +{:.2f}s, total {:.2f}s'.format(
num_ims, len(train_set.im_names),
time.time() - last_time,
time.time() - st_time))
last_time = time.time()
prob_diff = dict(zip(im_names, prob_diff))
return prob_diff
print('Sliding Window Occlusion for Non-mirrored Images.')
train_set.set_mirror_type(None)
prob_diff = dict()
prob_diff['non_mirrored'] = sw_occlude()
print('Sliding Window Occlusion for Mirrored Images.')
train_set.set_mirror_type('always')
prob_diff['mirrored'] = sw_occlude()
save_pickle(prob_diff, osp.join(exp_dir, 'prob_diff.pkl'))
save_pickle(all_masks, osp.join(exp_dir, 'all_masks.pkl'))