def predict_imgs(model, img_folder, bbox_folder, output_file, normalize,
detection_thresh):
detections = {}
for file in os.listdir(bbox_folder):
dets = load(os.path.join(bbox_folder, file))
assert dets.shape[1] == 5
img_name = file[:-4] # remove extension
detections[img_name] = dets
valid_dataset = hrnet_dataset.ImgFolderDataset(cfg, img_folder, detections,
normalize, detection_thresh)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=cfg.TEST.BATCH_SIZE_PER_GPU * len(cfg.GPUS),
shuffle=False,
pin_memory=True)
start = time.time()
preds, boxes, image_names, orig_boxes = predict(cfg, valid_loader,
valid_dataset, model)
end = time.time()
print("Time in prediction: " + str(end - start))
ensuredir(os.path.dirname(output_file))
valid_dataset.rescore_and_save_result(output_file, preds, boxes,
image_names, orig_boxes)
def main(args):
assert args.output.endswith('.pkl'), "Output file must be a pkl file"
# Clear up temp folder if it exists
if os.path.exists(args.tmp_folder) and not args.tpn_only:
shutil.rmtree(args.tmp_folder)
frame_dir = os.path.join(args.tmp_folder, 'frames')
bbox_dir = os.path.join(args.tmp_folder, 'bboxes')
keypoint_file = os.path.join(args.tmp_folder, 'keypoints.json')
if not args.tpn_only:
# split to frames:
ensuredir(os.path.join(args.tmp_folder, 'frames'))
out = os.system("ffmpeg -i %s -qscale:v 2 %s/frames/img_%%06d.jpg" %
(args.vid_path, args.tmp_folder))
if out != 0:
print("could not split to frames, error code: " + str(out))
sys.exit(1)
# Mask-RCNN
maskrcnn_bboxes.predict_imgs(frame_dir, bbox_dir)
# hrnet
hrnet_predict.predict(
'../hrnet/experiments/coco/hrnet/w32_256x192_adam_lr1e-3.yaml',
frame_dir, bbox_dir, keypoint_file)
# Run TPN
print("Running TPN...")
poses = run_tpn(args.model, frame_dir, keypoint_file, args.pose_refine,
args.focal_length, args.cx, args.cy)
save(args.output, poses)
def generate_vid_frames(cam, vid_id):
print(cam, vid_id)
metas = sequence_metas[cam][vid_id]
steps = [
2 if mpii_3dhp.get_train_fps(meta[0], meta[1]) == 50 else 1
for meta in metas
]
out_folder = os.path.join(muco_temp.MUCO_TEMP_PATH,
'frames/cam_%d/vid_%d' % (cam, vid_id))
ensuredir(out_folder)
gt_poses = load(
os.path.join(muco_temp.MUCO_TEMP_PATH,
'frames/cam_%d/gt.pkl' % cam))[vid_id]['annot3']
hip_ind = MuPoTSJoints().index_of('hip')
for i in range(NUM_FRAMES):
# generate frame
depths = gt_poses[i, :, hip_ind, 2]
ordered_poses = np.argsort(
depths)[::-1] # poses ordered by depth in decreasing order
bg_ind = ordered_poses[0]
img = mpii_3dhp.get_image(metas[bg_ind][0],
metas[bg_ind][1],
cam,
metas[bg_ind][2] + i * steps[bg_ind],
rgb=False)
img = img.astype('float32')
# add new pose onto image
for pose_ind in ordered_poses[1:]:
sub, seq, start = metas[pose_ind]
pose_img = mpii_3dhp.get_image(sub,
seq,
cam,
start + i * steps[pose_ind],
rgb=False)
# mask is 0 at greenscreen bg, 1 at foreground (body, chair)
mask = mpii_3dhp.get_mask(sub, seq, cam,
start + i * steps[pose_ind],
'FGmasks')[:, :, 2] / 255.
mask = cv2.GaussianBlur(mask, (0, 0), 2)[:, :, np.newaxis]
# chair_mask is 0 at chair, 1 everywhere else
chair_mask = mpii_3dhp.get_mask(sub, seq, cam,
start + i * steps[pose_ind],
'ChairMasks')[:, :, [2]] / 255
img = chair_mask * img + (1 - chair_mask) * pose_img
img = mask * pose_img + (1 - mask) * img
img = img.astype('uint8')
cv2.imwrite(os.path.join(out_folder, 'img_%04d.jpg' % i), img,
[cv2.IMWRITE_JPEG_QUALITY, 80])
def run_experiment(test_dict_2d, test_dict_3d, train_dict_2d, train_dict_3d,
train_dict_3d_absolute, test_dict_actionwise, cam_dict,
normalisation_params, out_path, p):
"""
Runs the experiment specified by the params object `p`. It saves additional data in a
timestamped folder under LOG_PATH, like network structure, trained model, training loss/errors, etc.
"""
model = siamese_network(p)
print "Parameters:"
print p
ensuredir(out_path)
save(os.path.join(out_path, 'params.txt'), str(p))
save(os.path.join(out_path, 'normalisation.pkl'), normalisation_params)
train_generator = PosNegBatchGenerator(
train_dict_2d,
train_dict_3d,
p,
shuffle=True,
cams=cam_dict,
train_dict_3d_absolute=train_dict_3d_absolute)
valid_generator = SplitBatchGenerator(test_dict_2d, test_dict_3d, p)
callbacks = []
if p.weight_decay:
callbacks.append(LearningRateScheduler(exp_decay(p), verbose=1))
callbacks.append(
LogAllMillimeterError(normalisation_params['std_3d'],
test_dict_actionwise,
siamese=len(model.inputs) >= 2,
csv=os.path.join(out_path, 'metrics.log')))
print "Output folder is", out_path
result = model.fit_generator(train_generator,
steps_per_epoch=p.TRAIN_SIZE / p.batch_size,
epochs=p.num_epochs,
callbacks=callbacks,
validation_data=valid_generator,
validation_steps=p.VALID_SIZE / p.batch_size,
verbose=2)
model.save(os.path.join(out_path, 'model.h5'))
return result
def predict_imgs(model, img_folder, bbox_folder, output_file, normalize,
detection_thresh):
detections = {}
for file in sorted(os.listdir(bbox_folder)):
dets = load(os.path.join(bbox_folder, file))
assert dets.shape[1] == 5
img_name = file[:-4] # remove extension
detections[img_name] = dets
valid_dataset = hrnet_dataset.ImgFolderDataset(cfg, img_folder, detections,
normalize, detection_thresh)
start = time.time()
preds, boxes, image_names, orig_boxes = predict_dataset(
cfg, valid_dataset, model)
end = time.time()
print("Time in prediction: " + str(end - start))
ensuredir(os.path.dirname(output_file))
valid_dataset.rescore_and_save_result(output_file, preds, boxes,
image_names, orig_boxes)
def run_experiment(output_path, _config, exp: Experiment):
exp.log_parameters(flatten_params(_config))
save(os.path.join(output_path, "config.json"), _config)
ensuredir(output_path)
if _config["train_data"] == "mpii_train":
print("Training data is mpii-train")
train_data = Mpi3dTrainDataset(
_config["pose2d_type"],
_config["pose3d_scaling"],
_config["cap_25fps"],
_config["stride"],
)
elif _config["train_data"] == "mpii+muco":
print("Training data is mpii-train and muco_temp concatenated")
mpi_data = Mpi3dTrainDataset(
_config["pose2d_type"],
_config["pose3d_scaling"],
_config["cap_25fps"],
_config["stride"],
)
muco_data = PersonStackedMucoTempDataset(_config["pose2d_type"],
_config["pose3d_scaling"])
train_data = ConcatPoseDataset(mpi_data, muco_data)
elif _config["train_data"].startswith("muco_temp"):
train_data = PersonStackedMucoTempDataset(_config["pose2d_type"],
_config["pose3d_scaling"])
test_data = Mpi3dTestDataset(_config["pose2d_type"],
_config["pose3d_scaling"],
eval_frames_only=True)
if _config["simple_aug"]:
train_data.augment(False)
assert _config["orient_norm"] == "gauss"
normalizer_orient = MeanNormalizeOrient(train_data)
# Load the preprocessing steps
train_data.transform = None
transforms_train = [
decode_trfrm(_config["preprocess_2d"], globals())(train_data,
cache=False),
decode_trfrm(_config["preprocess_3d"], globals())(train_data,
cache=False),
normalizer_orient,
]
normalizer2d = transforms_train[0].normalizer
normalizer3d = transforms_train[1].normalizer
transforms_test = [
decode_trfrm(_config["preprocess_2d"], globals())(test_data,
normalizer2d),
decode_trfrm(_config["preprocess_3d"], globals())(test_data,
normalizer3d),
normalizer_orient,
]
transforms_train.append(RemoveIndex())
transforms_test.append(RemoveIndex())
train_data.transform = SaveableCompose(transforms_train)
test_data.transform = SaveableCompose(transforms_test)
# save normalisation params
save(output_path + "/preprocess_params.pkl",
train_data.transform.state_dict())
len_train = len(train_data)
len_test = len(test_data)
print("Length of training data:", len_train)
print("Length of test data:", len_test)
exp.log_parameter("train data length", len_train)
exp.log_parameter("test data length", len_test)
bos = train_data[[0]]["orientation"].shape
out_shape = (bos[1] * bos[2] if _config["model"]["loss"] == "orient" else
MuPoTSJoints.NUM_JOINTS * 3)
model = TemporalModelOptimized1f(
train_data[[0]]["pose2d"].shape[-1],
out_shape,
_config["model"]["filter_widths"],
dropout=_config["model"]["dropout"],
channels=_config["model"]["channels"],
layernorm=_config["model"]["layernorm"],
)
test_model = TemporalModel(
train_data[[0]]["pose2d"].shape[-1],
out_shape,
_config["model"]["filter_widths"],
dropout=_config["model"]["dropout"],
channels=_config["model"]["channels"],
layernorm=_config["model"]["layernorm"],
)
model.cuda()
test_model.cuda()
save(output_path + "/model_summary.txt", str(model))
pad = (model.receptive_field() - 1) // 2
train_loader = ChunkedGenerator(
train_data,
_config["batch_size"],
pad,
_config["train_time_flip"],
shuffle=_config["shuffle"],
ordered_batch=_config["ordered_batch"],
)
tester = ModelCopyTemporalEvaluator(test_model,
test_data,
_config["model"]["loss"],
_config["test_time_flip"],
post_process3d=get_postprocessor(
_config, test_data, normalizer3d),
prefix="test",
orient_norm=_config["orient_norm"],
normalizer_orient=normalizer_orient)
torch_train(
exp,
train_loader,
model,
lambda m, b: calc_loss(m, b, _config, None, None, None,
normalizer_orient),
# lambda m, b: calc_loss(
# m,
# b,
# _config,
# torch.tensor(normalizer2d.mean[2::3]).cuda(),
# torch.tensor(normalizer2d.std[2::3]).cuda(),
# torch.tensor(normalizer3d.std).cuda(),
# ),
_config,
callbacks=[tester],
)
model_path = os.path.join(output_path, "model_params.pkl")
torch.save(model.state_dict(), model_path)
exp.log_model("model", model_path)
save(
output_path + "/test_results.pkl",
{
"index": test_data.index,
"pred": preds_from_logger(test_data, tester),
"pose3d": test_data.poses3d,
},
)
def run_experiment(output_path, _config, exp: Experiment):
config, m = eval.load_model(_config["weights"])
# config.update(_config)
config["model"].update(_config["model"])
_config["model"] = config["model"]
# tmp = _config["model"]["loss"]
# _config["model"]["loss"] = "v * mse + e_smooth_small"
exp.log_parameters(train.flatten_params(_config))
# _config["model"]["loss"] = tmp
save(os.path.join(output_path, "config.json"), _config)
ensuredir(output_path)
if _config["train_data"] == "mpii_train":
print("Training data is mpii-train")
train_data = Mpi3dTrainDataset(
_config["pose2d_type"],
_config["pose3d_scaling"],
_config["cap_25fps"],
_config["stride"],
)
elif _config["train_data"] == "mpii+muco":
print("Training data is mpii-train and muco_temp concatenated")
mpi_data = Mpi3dTrainDataset(
_config["pose2d_type"],
_config["pose3d_scaling"],
_config["cap_25fps"],
_config["stride"],
)
muco_data = PersonStackedMucoTempDataset(_config["pose2d_type"],
_config["pose3d_scaling"])
train_data = ConcatPoseDataset(mpi_data, muco_data)
elif _config["train_data"].startswith("muco_temp"):
train_data = PersonStackedMucoTempDataset(_config["pose2d_type"],
_config["pose3d_scaling"])
test_data = Mpi3dTestDataset(_config["pose2d_type"],
_config["pose3d_scaling"],
eval_frames_only=True)
if _config["simple_aug"]:
train_data.augment(False)
# Load the preprocessing steps
params_path = os.path.join(LOG_PATH, _config["weights"],
"preprocess_params.pkl")
transform = SaveableCompose.from_file(params_path, test_data, globals())
train_data.transform = None
transforms_train = [
decode_trfrm(_config["preprocess_2d"], globals())(train_data,
cache=False),
decode_trfrm(_config["preprocess_3d"], globals())(train_data,
cache=False),
]
normalizer2d = transforms_train[0].normalizer
normalizer3d = transforms_train[1].normalizer
transforms_test = [
decode_trfrm(_config["preprocess_2d"], globals())(test_data,
normalizer2d),
decode_trfrm(_config["preprocess_3d"], globals())(test_data,
normalizer3d),
]
transforms_train.append(RemoveIndex())
transforms_test.append(RemoveIndex())
train_data.transform = SaveableCompose(transforms_train)
test_data.transform = SaveableCompose(transforms_test)
# train_data.transform = SaveableCompose.from_file(params_path, train_data, globals())
# test_data.transform = SaveableCompose.from_file(params_path, test_data, globals())
# save normalisation params
save(output_path + "/preprocess_params.pkl",
train_data.transform.state_dict())
len_train = len(train_data)
len_test = len(test_data)
print("Length of training data:", len_train)
print("Length of test data:", len_test)
exp.log_parameter("train data length", len_train)
exp.log_parameter("test data length", len_test)
model = TemporalModelOptimized1f(
train_data[[0]]["pose2d"].shape[-1],
MuPoTSJoints.NUM_JOINTS,
config["model"]["filter_widths"],
dropout=config["model"]["dropout"],
channels=config["model"]["channels"],
layernorm=config["model"]["layernorm"],
)
model.load_state_dict(m.state_dict())
test_model = TemporalModel(
train_data[[0]]["pose2d"].shape[-1],
MuPoTSJoints.NUM_JOINTS,
config["model"]["filter_widths"],
dropout=config["model"]["dropout"],
channels=config["model"]["channels"],
layernorm=config["model"]["layernorm"],
)
model.cuda()
test_model.cuda()
save(output_path + "/model_summary.txt", str(model))
# normalizer2d = train_data.transform.transforms[0].normalizer
# normalizer3d = train_data.transform.transforms[1].normalizer
pad = (model.receptive_field() - 1) // 2
train_loader = ChunkedGenerator(
train_data,
_config["batch_size"],
pad,
_config["train_time_flip"],
shuffle=_config["shuffle"],
ordered_batch=_config["ordered_batch"],
)
tester = ModelCopyTemporalEvaluator(
test_model,
test_data,
config["model"]["loss"],
_config["test_time_flip"],
post_process3d=get_postprocessor(_config, test_data, normalizer3d),
prefix="test",
)
torch_train(
exp,
train_loader,
model,
lambda m, b: train.calc_loss(
m,
b,
_config,
torch.tensor(normalizer2d.mean[2::3]).cuda(),
torch.tensor(normalizer2d.std[2::3]).cuda(),
torch.tensor(normalizer3d.std).cuda(),
),
_config,
callbacks=[tester],
)
model_path = os.path.join(output_path, "model_params.pkl")
torch.save(model.state_dict(), model_path)
exp.log_model("model", model_path)
def rescore_and_save_result(self, output_file, preds, all_boxes, img_path,
orig_boxes):
assert output_file.endswith('.json') or output_file.endswith(
'.npy'), "Only json and numpy output is supported"
ensuredir(os.path.dirname(output_file))
# person x (keypoints)
_kpts = []
for idx, kpt in enumerate(preds):
_kpts.append({
'keypoints': kpt,
'center': all_boxes[idx][0:2],
'scale': all_boxes[idx][2:4],
'area': all_boxes[idx][4],
'score': all_boxes[idx][5],
'image': img_path[idx],
'origbox': orig_boxes[idx]
})
# image x person x (keypoints)
kpts = defaultdict(list)
for kpt in _kpts:
kpts[kpt['image']].append(kpt)
# rescoring and oks nms
num_joints = self.num_joints
in_vis_thre = self.in_vis_thre
oks_thre = self.oks_thre
oks_nmsed_kpts = []
nmsed_kpts_by_frame = defaultdict(list)
for img in kpts.keys():
img_kpts = kpts[img]
for n_p in img_kpts:
box_score = n_p['score']
kpt_score = 0
valid_num = 0
for n_jt in range(0, num_joints):
t_s = n_p['keypoints'][n_jt][2]
if t_s > in_vis_thre:
kpt_score = kpt_score + t_s
valid_num = valid_num + 1
if valid_num != 0:
kpt_score = kpt_score / valid_num
# rescoring
n_p['score'] = kpt_score * box_score
if self.soft_nms:
keep = soft_oks_nms(
[img_kpts[i] for i in range(len(img_kpts))], oks_thre)
else:
keep = oks_nms([img_kpts[i] for i in range(len(img_kpts))],
oks_thre)
if len(keep) == 0:
selected_kpts = img_kpts
else:
selected_kpts = [img_kpts[_keep] for _keep in keep]
oks_nmsed_kpts.append(selected_kpts)
nmsed_kpts_by_frame[img] = selected_kpts
self._write_keypoint_results(nmsed_kpts_by_frame, output_file)
def run_experiment(output_path, _config):
save(os.path.join(output_path, 'config.json'), _config)
ensuredir(output_path)
if _config['train_data'] == 'mpii_train':
print("Training data is mpii-train")
train_data = Mpi3dTrainDataset(_config['pose2d_type'],
_config['pose3d_scaling'],
_config['cap_25fps'], _config['stride'])
elif _config['train_data'] == 'mpii+muco':
print("Training data is mpii-train and muco_temp concatenated")
mpi_data = Mpi3dTrainDataset(_config['pose2d_type'],
_config['pose3d_scaling'],
_config['cap_25fps'], _config['stride'])
muco_data = PersonStackedMucoTempDataset(_config['pose2d_type'],
_config['pose3d_scaling'])
train_data = ConcatPoseDataset(mpi_data, muco_data)
elif _config['train_data'].startswith('muco_temp'):
train_data = PersonStackedMucoTempDataset(_config['pose2d_type'],
_config['pose3d_scaling'])
test_data = Mpi3dTestDataset(_config['pose2d_type'],
_config['pose3d_scaling'],
eval_frames_only=True)
if _config['simple_aug']:
train_data.augment(False)
# Load the preprocessing steps
train_data.transform = None
transforms_train = [
decode_trfrm(_config['preprocess_2d'], globals())(train_data,
cache=False),
decode_trfrm(_config['preprocess_3d'], globals())(train_data,
cache=False)
]
normalizer2d = transforms_train[0].normalizer
normalizer3d = transforms_train[1].normalizer
transforms_test = [
decode_trfrm(_config['preprocess_2d'], globals())(test_data,
normalizer2d),
decode_trfrm(_config['preprocess_3d'], globals())(test_data,
normalizer3d)
]
transforms_train.append(RemoveIndex())
transforms_test.append(RemoveIndex())
train_data.transform = SaveableCompose(transforms_train)
test_data.transform = SaveableCompose(transforms_test)
# save normalisation params
save(output_path + '/preprocess_params.pkl',
train_data.transform.state_dict())
print("Length of training data:", len(train_data))
print("Length of test data:", len(test_data))
model = TemporalModelOptimized1f(train_data[[0]]['pose2d'].shape[-1],
MuPoTSJoints.NUM_JOINTS,
_config['model']['filter_widths'],
dropout=_config['model']['dropout'],
channels=_config['model']['channels'],
layernorm=_config['model']['layernorm'])
test_model = TemporalModel(train_data[[0]]['pose2d'].shape[-1],
MuPoTSJoints.NUM_JOINTS,
_config['model']['filter_widths'],
dropout=_config['model']['dropout'],
channels=_config['model']['channels'],
layernorm=_config['model']['layernorm'])
model.cuda()
test_model.cuda()
save(output_path + '/model_summary.txt', str(model))
pad = (model.receptive_field() - 1) // 2
train_loader = ChunkedGenerator(train_data,
_config['batch_size'],
pad,
_config['train_time_flip'],
shuffle=True)
tester = ModelCopyTemporalEvaluator(test_model,
test_data,
_config['model']['loss'],
_config['test_time_flip'],
post_process3d=get_postprocessor(
_config, test_data, normalizer3d),
prefix='test')
torch_train(train_loader,
model,
lambda m, b: calc_loss(m, b, _config),
_config,
callbacks=[tester])
torch.save(model.state_dict(), os.path.join(output_path,
'model_params.pkl'))
save(
output_path + '/test_results.pkl', {
'index': test_data.index,
'pred': preds_from_logger(test_data, tester),
'pose3d': test_data.poses3d
})