Spaces:
Runtime error
Runtime error
File size: 9,625 Bytes
753fd9a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 |
# scripts/train.py --workers 12 --checkpoint project22_no3dcgloss_smaldogsilvia_v0 --loss-weight-path barc_loss_weights_no3dcgloss.json --config barc_cfg_train.yaml start --model-file-hg hg_ksp_fromnewanipose_stanext_v0/checkpoint.pth.tar --model-file-3d barc_normflow_pret/checkpoint.pth.tar
import torch
import torch.backends.cudnn
import torch.nn.parallel
from tqdm import tqdm
import os
import json
import pathlib
import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../', 'src'))
# from stacked_hourglass.loss import joints_mse_loss
from stacked_hourglass.loss import joints_mse_loss_onKPloc
from stacked_hourglass.utils.evaluation import accuracy, AverageMeter, final_preds, get_preds, get_preds_soft
from stacked_hourglass.utils.transforms import fliplr, flip_back
from stacked_hourglass.utils.visualization import save_input_image_with_keypoints
def do_training_step(model, optimiser, input, target, meta, data_info, target_weight=None):
assert model.training, 'model must be in training mode.'
assert len(input) == len(target), 'input and target must contain the same number of examples.'
with torch.enable_grad():
# Forward pass and loss calculation.
output = model(input)
# original: loss = sum(joints_mse_loss(o, target, target_weight) for o in output)
# NEW:
loss = sum(joints_mse_loss_onKPloc(o, target, meta, target_weight) for o in output)
# Backward pass and parameter update.
optimiser.zero_grad()
loss.backward()
optimiser.step()
return output[-1], loss.item()
def do_training_epoch(train_loader, model, device, data_info, optimiser, quiet=False, acc_joints=None):
losses = AverageMeter()
accuracies = AverageMeter()
# Put the model in training mode.
model.train()
iterable = enumerate(train_loader)
progress = None
if not quiet:
progress = tqdm(iterable, desc='Train', total=len(train_loader), ascii=True, leave=False)
iterable = progress
for i, (input, target, meta) in iterable:
input, target = input.to(device), target.to(device, non_blocking=True)
target_weight = meta['target_weight'].to(device, non_blocking=True)
output, loss = do_training_step(model, optimiser, input, target, meta, data_info, target_weight)
acc = accuracy(output, target, acc_joints)
# measure accuracy and record loss
losses.update(loss, input.size(0))
accuracies.update(acc[0], input.size(0))
# Show accuracy and loss as part of the progress bar.
if progress is not None:
progress.set_postfix_str('Loss: {loss:0.4f}, Acc: {acc:6.2f}'.format(
loss=losses.avg,
acc=100 * accuracies.avg
))
return losses.avg, accuracies.avg
def do_validation_step(model, input, target, meta, data_info, target_weight=None, flip=False):
# assert not model.training, 'model must be in evaluation mode.'
assert len(input) == len(target), 'input and target must contain the same number of examples.'
# Forward pass and loss calculation.
output = model(input)
# original: loss = sum(joints_mse_loss(o, target, target_weight) for o in output)
# NEW:
loss = sum(joints_mse_loss_onKPloc(o, target, meta, target_weight) for o in output)
# Get the heatmaps.
if flip:
# If `flip` is true, perform horizontally flipped inference as well. This should
# result in more robust predictions at the expense of additional compute.
flip_input = fliplr(input)
flip_output = model(flip_input)
flip_output = flip_output[-1].cpu()
flip_output = flip_back(flip_output.detach(), data_info.hflip_indices)
heatmaps = (output[-1].cpu() + flip_output) / 2
else:
heatmaps = output[-1].cpu()
return heatmaps, loss.item()
def do_validation_epoch(val_loader, model, device, data_info, flip=False, quiet=False, acc_joints=None, save_imgs_path=None):
losses = AverageMeter()
accuracies = AverageMeter()
predictions = [None] * len(val_loader.dataset)
if save_imgs_path is not None:
pathlib.Path(save_imgs_path).mkdir(parents=True, exist_ok=True)
# Put the model in evaluation mode.
model.eval()
iterable = enumerate(val_loader)
progress = None
if not quiet:
progress = tqdm(iterable, desc='Valid', total=len(val_loader), ascii=True, leave=False)
iterable = progress
for i, (input, target, meta) in iterable:
# Copy data to the training device (eg GPU).
input = input.to(device, non_blocking=True)
target = target.to(device, non_blocking=True)
target_weight = meta['target_weight'].to(device, non_blocking=True)
# import pdb; pdb.set_trace()
heatmaps, loss = do_validation_step(model, input, target, meta, data_info, target_weight, flip)
# Calculate PCK from the predicted heatmaps.
acc = accuracy(heatmaps, target.cpu(), acc_joints)
# Calculate locations in original image space from the predicted heatmaps.
preds = final_preds(heatmaps, meta['center'], meta['scale'], [64, 64])
# NEW for visualization: (and redundant, but for visualization)
preds_unprocessed, preds_unprocessed_maxval = get_preds_soft(heatmaps, return_maxval=True)
# preds_unprocessed, preds_unprocessed_norm, preds_unprocessed_maxval = get_preds_soft(heatmaps, return_maxval=True, norm_and_unnorm_coords=True)
# import pdb; pdb.set_trace()
ind = 0
for example_index, pose in zip(meta['index'], preds):
predictions[example_index] = pose
# NEW for visualization
if save_imgs_path is not None:
out_name = os.path.join(save_imgs_path, 'res_' + str( example_index.item()) + '.png')
pred_unp = preds_unprocessed[ind, :, :]
pred_unp_maxval = preds_unprocessed_maxval[ind, :, :]
pred_unp_prep = torch.cat((pred_unp, pred_unp_maxval), 1)
inp_img = input[ind, :, :, :]
# the following line (with -1) should not be needed anymore after cvpr (after bugfix01 in data preparation 08.09.2022)
# pred_unp_prep[:, :2] = pred_unp_prep[:, :2] - 1
# save_input_image_with_keypoints(inp_img, pred_unp_prep, out_path=out_name, threshold=0.1, print_scores=True) # here we have default ratio_in_out=4.
# NEW: 08.09.2022 after bugfix01
# import pdb; pdb.set_trace()
pred_unp_prep[:, :2] = pred_unp_prep[:, :2] * 4
if 'name' in meta.keys(): # we do this for the stanext set
name = meta['name'][ind]
out_path_keyp_img = os.path.join(os.path.dirname(out_name), name)
out_path_json = os.path.join(os.path.dirname(out_name), name).replace('_vis', '_json').replace('.jpg', '.json')
if not os.path.exists(os.path.dirname(out_path_json)):
os.makedirs(os.path.dirname(out_path_json))
if not os.path.exists(os.path.dirname(out_path_keyp_img)):
os.makedirs(os.path.dirname(out_path_keyp_img))
save_input_image_with_keypoints(inp_img, pred_unp_prep, out_path=out_path_keyp_img, ratio_in_out=1.0, threshold=0.1, print_scores=True) # threshold=0.3
out_name_json = out_path_json # os.path.join(save_imgs_path, 'res_' + str( example_index.item()) + '.json')
res_dict = {
'pred_joints_256': list(pred_unp_prep.cpu().numpy().astype(float).reshape((-1))),
'center': list(meta['center'][ind, :].cpu().numpy().astype(float).reshape((-1))),
'scale': meta['scale'][ind].item()}
with open(out_name_json, 'w') as outfile: json.dump(res_dict, outfile)
else:
save_input_image_with_keypoints(inp_img, pred_unp_prep, out_path=out_name, ratio_in_out=1.0, threshold=0.1, print_scores=True) # threshold=0.3
'''# animalpose_hg8_v0 (did forget to subtract 1 in dataset)
pred_unp_prep[:, :2] = pred_unp_prep[:, :2] * 4 ############ Why is this necessary???
pred_unp_prep[:, :2] = pred_unp_prep[:, :2] - 1
save_input_image_with_keypoints(inp_img, pred_unp_prep, out_path=out_name, ratio_in_out=1.0, threshold=0.1, print_scores=True) # threshold=0.3
out_name_json = os.path.join(save_imgs_path, 'res_' + str( example_index.item()) + '.json')
res_dict = {
'pred_joints_256': list(pred_unp_prep.cpu().numpy().astype(float).reshape((-1))),
'center': list(meta['center'][ind, :].cpu().numpy().astype(float).reshape((-1))),
'scale': meta['scale'][ind].item()}
with open(out_name_json, 'w') as outfile: json.dump(res_dict, outfile)'''
ind += 1
# Record accuracy and loss for this batch.
losses.update(loss, input.size(0))
accuracies.update(acc[0].item(), input.size(0))
# Show accuracy and loss as part of the progress bar.
if progress is not None:
progress.set_postfix_str('Loss: {loss:0.4f}, Acc: {acc:6.2f}'.format(
loss=losses.avg,
acc=100 * accuracies.avg
))
predictions = torch.stack(predictions, dim=0)
return losses.avg, accuracies.avg, predictions
|