def validate(val_loader, epoch, model, criterion, args, save_path=False):
# keeping statistics
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
mse = AverageMeter()
mse_hip = AverageMeter()
mse_coxa = AverageMeter()
mse_femur = AverageMeter()
mse_tibia = AverageMeter()
mse_tarsus = AverageMeter()
mse_jump = AverageMeter()
avg_local_max = AverageMeter()
# predictions and score maps
num_cameras = 7
predictions = np.zeros(
shape=(
num_cameras + 1,
val_loader.dataset.greatest_image_id() + 1,
config["num_predict"],
2,
),
dtype=np.float32,
) # num_cameras+1 for the mirrored camera 3
getLogger('df3d').debug("Predictions shape {}".format(predictions.shape))
if save_path is not None:
print("SAVE_PATH:")
print(save_path)
unlabeled_folder_replace = save_path.replace("/", "-")
# score_map_filename = os.path.join(args.data_folder, "./heatmap_{}.pkl".format(unlabeled_folder_replace))
score_map_filename = os.path.join(
args.data_folder,
"{}".format(save_path),
args.output_folder,
"heatmap_{}.pkl".format(unlabeled_folder_replace),
)
score_map_path = Path(score_map_filename)
score_map_path.parent.mkdir(exist_ok=True, parents=True)
score_map_arr = np.memmap(
score_map_filename,
dtype="float32",
mode="w+",
shape=(
num_cameras + 1,
val_loader.dataset.greatest_image_id() + 1,
config["num_predict"],
args.hm_res[0],
args.hm_res[1],
),
) # num_cameras+1 for the mirrored camera 3
# switch to evaluate mode
model.eval()
end = time.time()
bar = Bar("Processing", max=len(val_loader)) #if logging.getLogger('df3d').isEnabledFor(logging.INFO) else NoOutputBar()
bar.start()
for i, (inputs, target, meta) in enumerate(val_loader):
# measure data loading time
data_time.update(time.time() - end)
target = target.cuda(non_blocking=True)
input_var = torch.autograd.Variable(inputs.cuda())
target_var = torch.autograd.Variable(target)
# compute output
with torch.no_grad():
output = model(input_var)
score_map = output[-1].data.cpu()
# loss = criterion(output[-1], target_var)
acc = accuracy(score_map.cpu(), target.cpu(), args.acc_joints)
# generate predictions
if save_path is not None:
for n, cam_id, cam_read_id, img_id in zip(
range(score_map.size(0)), meta["cid"], meta["cam_read_id"], meta["pid"]
):
smap = to_numpy(score_map[n, :, :, :])
score_map_arr[cam_read_id, img_id, :] = smap
pr = Camera.hm_to_pred(
score_map_arr[cam_read_id, img_id, :], threshold_abs=0.0
)
predictions[cam_read_id, img_id, :] = pr
# measure accuracy and record loss
mse_err = mse_acc(target_var.data.cpu(), score_map)
mse.update(torch.mean(mse_err[args.acc_joints, :]), inputs.size(0))
mse_hip.update(
torch.mean(mse_err[np.arange(0, mse_err.shape[0], 5), :]), inputs.size(0)
)
mse_coxa.update(
torch.mean(mse_err[np.arange(1, mse_err.shape[0], 5), :]), inputs.size(0)
)
mse_femur.update(
torch.mean(mse_err[np.arange(2, mse_err.shape[0], 5), :]), inputs.size(0)
)
mse_tibia.update(
torch.mean(mse_err[np.arange(3, mse_err.shape[0], 5), :]), inputs.size(0)
)
mse_tarsus.update(
torch.mean(mse_err[np.arange(4, mse_err.shape[0], 5), :]), inputs.size(0)
)
losses.update(0, inputs.size(0))
acces.update(acc[0], inputs.size(0))
jump_thr = args.hm_res[0] * 5.0 / 64.0
mse_jump.update(
np.mean(np.ravel(mse_err[args.acc_joints, :] > jump_thr)) * 100.0,
inputs.size(0),
)
if args.unlabeled and i < args.num_output_image:
drosophila_img = drosophila_image_overlay(
inputs,
score_map,
args.hm_res,
3,
np.arange(config["num_predict"]),
img_id=int(meta["image_name"][0].split("_")[-1]),
)
template = "{}_overlay_val_epoch_{}_{}_{:06d}_{}.jpg"
folder_name = meta["folder_name"][0]
folder_name = (
folder_name.replace(" ", "_").replace("\\", "-").replace("/", "-")
)
save_image_name = template.format(
"left" if meta["cid"][0] == 0 else "right",
epoch,
folder_name,
meta["pid"][0],
meta["cid"][0],
)
save_image_name = save_image_name.replace(" ", "_ ").replace("\\", "-")
save_image(
img_path=os.path.join(args.checkpoint_image_dir, save_image_name),
img=drosophila_img,
)
elif i < args.num_output_image:
drosophila_img = drosophila_image_overlay(
inputs, score_map, args.hm_res, 3, np.arange(config["num_predict"])
)
folder_name = meta["folder_name"][0]
template = "./overlay_val_epoch_{}_{}_{:06d}_{}.jpg"
save_image_name = template.format(
epoch, folder_name.replace("/", "-"), meta["pid"][0], meta["cid"][0]
)
save_image_name = save_image_name.replace(" ", "_ ").replace("\\", "-")
save_image(
img_path=os.path.join(args.checkpoint_image_dir, save_image_name),
img=drosophila_img,
)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = "({batch}/{size}) D:{data:.6f}s|B:{bt:.3f}s|L:{loss:.8f}|Acc:{acc: .4f}|Mse:{mse: .3f}|F-T:{mse_femur: .3f}|T-Tar:{mse_tibia: .3f}|Tar-tip:{mse_tarsus: .3f}".format(
batch=i + 1,
size=len(val_loader),
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg,
mse=mse.avg,
mse_hip=mse_hip.avg,
mse_femur=mse_femur.avg,
mse_tibia=mse_tibia.avg,
mse_tarsus=mse_tarsus.avg,
mse_coxa=mse_coxa.avg,
mse_jump=mse_jump.avg,
local_max_count=avg_local_max.avg,
)
bar.next()
bar.finish()
if save_path is not None:
score_map_arr.flush()
else:
score_map_arr = None
return losses.avg, acces.avg, predictions, score_map_arr, mse.avg, mse_jump.avg
def step(loader, model, optimizer, mode, heatmap, epoch, num_classes,
acc_joints):
keys_am = [
"batch_time",
"data_time",
"losses",
"acces",
"mse",
"body_coxa",
"coxa_femur",
"femur_tibia",
"tibia_tarsus",
"tarsus_tip",
]
am = {k: AverageMeter() for k in keys_am}
if mode == mode.train:
model.train()
elif mode == mode.test:
model.eval()
start = time.time()
bar = (Bar("Processing", max=len(loader))
if logger.debug_enabled() else NoOutputBar())
bar.start()
predictions = np.zeros(
shape=(
config["num_cameras"] + 1,
loader.dataset.greatest_image_id() + 1,
config["num_predict"],
2,
),
dtype=np.float32,
) # num_cameras+1 for the mirrored camera 3
for i, (inputs, target, meta) in enumerate(loader):
np.random.seed()
am["data_time"].update(time.time() - start)
input_var = torch.autograd.Variable(on_cuda(inputs))
target_var = torch.autograd.Variable(on_cuda(target,
non_blocking=True))
output = model(input_var)
heatmap_batch = output[-1].data.cpu()
for n, cam_read_id, img_id in zip(range(heatmap_batch.size(0)),
meta["cam_read_id"], meta["pid"]):
smap = to_numpy(heatmap_batch[n, :, :, :])
pr = Camera.hm_to_pred(smap, threshold_abs=0.0)
predictions[cam_read_id, img_id, :] = pr
if heatmap is not None:
heatmap[cam_read_id, img_id, :] = smap
# loss = df3dLoss(output, target_var, meta["joint_exists"], num_classes)
if mode == mode.train:
optimizer.zero_grad()
loss.backward()
optimizer.step()
acc = accuracy(heatmap_batch, target, acc_joints)
# am["acces"].add(acc)
mse_err = mse_acc(target_var.data.cpu(), heatmap_batch)
# am["mse"].add(mse_err)
am["batch_time"].update(time.time() - start)
start = time.time()
bar.suffix = "{epoch} | ({batch}/{size}) D:{data:.6f}s|B:{bt:.3f}s|L:{loss:.8f}|Acc:{acc: .4f}|Mse:{mse: .3f} F-T:{mse_femur: .3f} T-Tar:{mse_tibia: .3f} Tar-tip:{mse_tarsus: .3f}".format(
epoch=epoch,
batch=i + 1,
size=len(loader),
data=am["data_time"].val,
bt=am["batch_time"].val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=am["losses"].avg,
acc=am["acces"].avg,
mse=am["mse"].avg,
mse_hip=am["body_coxa"].avg,
mse_femur=am["coxa_femur"].avg,
mse_tibia=am["femur_tibia"].avg,
mse_tarsus=am["tarsus_tip"].avg,
)
bar.next()
bar.finish()
return (predictions, heatmap, am["losses"].avg, am["acces"].avg,
am["mse"].avg)
def train(train_loader, epoch, model, optimizer, criterion, args):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
mse = AverageMeter()
mse_hip = AverageMeter()
mse_coxa = AverageMeter()
mse_femur = AverageMeter()
mse_tibia = AverageMeter()
mse_tarsus = AverageMeter()
mse_jump = AverageMeter()
avg_local_max = AverageMeter()
# to later save the disk
predictions = dict()
# switch to train mode
model.train()
end = time.time()
bar = Bar("Processing", max=len(train_loader))# if logging.getLogger('df3d').isEnabledFor(logging.DEBUG) else NoOutputBar()
bar.start()
for i, (inputs, target, meta) in enumerate(train_loader):
# reset seed for truly random transformations on input
np.random.seed()
# measure data loading time
data_time.update(time.time() - end)
input_var = (
torch.autograd.Variable(inputs.cuda())
if torch.cuda.is_available()
else torch.autograd.Variable(inputs)
)
target_var = torch.autograd.Variable(target.cuda(non_blocking=True))
# compute output
output = model(input_var)
score_map = output[-1].data.cpu()
loss_weight = torch.ones((inputs.size(0), args.num_classes, 1, 1))
if np.any(np.array(meta["joint_exists"]) == 0):
batch_index = (np.logical_not(meta["joint_exists"])).nonzero()[:, 0]
joint_index = (np.logical_not(meta["joint_exists"])).nonzero()[:, 1]
loss_weight[batch_index, joint_index, :, :] = 0.0
# logger.debug(loss_weight)
loss_weight.requires_grad = True
loss_weight = loss_weight.cuda()
loss = weighted_mse_loss(output[0], target_var, weights=loss_weight)
for j in range(1, len(output)):
loss += weighted_mse_loss(output[j], target_var, weights=loss_weight)
# Compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
acc = accuracy(score_map, target, args.acc_joints)
# measure accuracy and record loss
mse_err = mse_acc(target_var.data.cpu(), score_map)
getLogger('df3d').debug(f'{mse_err.size()}, {mse_err[:, 0] > 50}, {meta["image_name"][0]}')
mse.update(torch.mean(mse_err[args.acc_joints, :]), inputs.size(0)) # per joint
mse_hip.update(torch.mean(mse_err[np.arange(0, 15, 5), :]), inputs.size(0))
mse_coxa.update(torch.mean(mse_err[np.arange(1, 15, 5), :]), inputs.size(0))
mse_femur.update(torch.mean(mse_err[np.arange(2, 15, 5), :]), inputs.size(0))
mse_tibia.update(torch.mean(mse_err[np.arange(3, 15, 5), :]), inputs.size(0))
mse_tarsus.update(torch.mean(mse_err[np.arange(4, 15, 5), :]), inputs.size(0))
losses.update(loss.item(), inputs.size(0))
acces.update(acc[0], inputs.size(0))
jump_thr = args.hm_res[0] * 5.0 / 64.0
mse_jump.update(
np.mean(np.ravel(mse_err[args.acc_joints, :] > jump_thr)) * 100.0,
inputs.size(0),
)
if i < args.num_output_image:
drosophila_img = drosophila_image_overlay(
inputs, target, args.hm_res, 3, args.train_joints
)
folder_name = (
meta["folder_name"][0].replace("/", "-")
+ meta["folder_name"][1].replace("/", "-")
+ meta["folder_name"][2].replace("/", "-")
)
image_name = meta["image_name"][0]
template = "./train_gt_{}_{}_{:06d}_{}_joint_{}.jpg"
save_image(
img_path=os.path.join(
args.checkpoint_image_dir,
template.format(
epoch, folder_name, meta["pid"][0], meta["cid"][0], []
),
),
img=drosophila_img,
)
if i < args.num_output_image and False:
drosophila_img = drosophila_image_overlay(
inputs,
score_map,
args.hm_res,
3,
args.train_joints,
img_id=int(meta["image_name"][0].split("_")[-1]),
)
template = "{}_train_pred_epoch_{}_{}_{:06d}_{}.jpg"
folder_name = meta["folder_name"][0]
folder_name = (
folder_name.replace(" ", "_").replace("\\", "-").replace("/", "-")
)
save_image_name = template.format(
"left" if meta["cid"][0] == 0 else "right",
epoch,
folder_name,
meta["pid"][0],
meta["cid"][0],
)
save_image_name = save_image_name.replace(" ", "_ ").replace("\\", "-")
save_image(
img_path=os.path.join(args.checkpoint_image_dir, save_image_name),
img=drosophila_img,
)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = "({batch}/{size}) D:{data:.6f}s|B:{bt:.3f}s|L:{loss:.8f}|Acc:{acc: .4f}|Mse:{mse: .3f} F-T:{mse_femur: .3f} T-Tar:{mse_tibia: .3f} Tar-tip:{mse_tarsus: .3f} Jump:{mse_jump: .4f}".format(
batch=i + 1,
size=len(train_loader),
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg,
mse=mse.avg,
mse_hip=mse_hip.avg,
mse_femur=mse_femur.avg,
mse_tibia=mse_tibia.avg,
mse_tarsus=mse_tarsus.avg,
mse_coxa=mse_coxa.avg,
mse_jump=mse_jump.avg,
)
bar.next()
bar.finish()
return losses.avg, acces.avg, predictions, mse.avg, mse_jump.avg