def main():
# Parse the options
opts = Opts().parse()
opts.device = torch.device(f'cuda:{opts.gpu[0]}')
print(opts.expID, opts.task)
# Record the start time
time_start = time.time()
# TODO: select the dataset by the options
# Set up dataset
train_loader_unit = PENN_CROP(opts, 'train')
train_loader = tud.DataLoader(train_loader_unit,
batch_size=opts.trainBatch,
shuffle=False,
num_workers=int(opts.num_workers))
val_loader = tud.DataLoader(PENN_CROP(opts, 'val'),
batch_size=1,
shuffle=False,
num_workers=int(opts.num_workers))
# Read number of joints(dim of output) from dataset
opts.nJoints = train_loader_unit.part.shape[1]
# Create the Model, Optimizer and Criterion
if opts.loadModel == 'none':
model = Hourglass2DPrediction(opts).cuda(device=opts.device)
else:
model = torch.load(opts.loadModel).cuda(device=opts.device)
# Set the Criterion and Optimizer
criterion = torch.nn.MSELoss(reduce=False).cuda(device=opts.device)
# opts.nOutput = len(model.outnode.children)
optimizer = torch.optim.RMSprop(model.parameters(),
opts.LR,
alpha=opts.alpha,
eps=opts.epsilon,
weight_decay=opts.weightDecay,
momentum=opts.momentum)
# If TEST, just validate
# TODO: save the validate results to mat or hdf5
if opts.test:
loss_test, pck_test = val(0, opts, val_loader, model, criterion)
print(f"test: | loss_test: {loss_test}| PCK_val: {pck_test}\n")
## TODO: save the predictions for the test
#sio.savemat(os.path.join(opts.saveDir, 'preds.mat'), mdict = {'preds':preds})
return
# NOT TEST, Train and Validate
for epoch in range(1, opts.nEpochs + 1):
## Train the model
loss_train, pck_train = train(epoch, opts, train_loader, model,
criterion, optimizer)
## Show results and elapsed time
time_elapsed = time.time() - time_start
print(
f"epoch: {epoch} | loss_train: {loss_train} | PCK_train: {pck_train} | {time_elapsed//60:.0f}min {time_elapsed%60:.0f}s\n"
)
## Intervals to show eval results
if epoch % opts.valIntervals == 0:
# TODO: Test the validation part
### Validation
loss_val, pck_val = val(epoch, opts, val_loader, model, criterion)
print(
f"epoch: {epoch} | loss_val: {loss_val}| PCK_val: {pck_val}\n")
### Save the model
torch.save(model, os.path.join(opts.save_dir,
f"model_{epoch}.pth"))
### TODO: save the preds for the validation
#sio.savemat(os.path.join(opts.saveDir, f"preds_{epoch}.mat"), mdict={'preds':preds})
# Use the optimizer to adjust learning rate
if epoch % opts.dropLR == 0:
lr = adjust_learning_rate(optimizer, epoch, opts.dropLR, opts.LR)
print(f"Drop LR to {lr}\n")
def main():
# Parse the options from parameters
opts = Opts().parse()
## For PyTorch 0.4.1, cuda(device)
opts.device = torch.device(f'cuda:{opts.gpu[0]}')
print(opts.expID, opts.task, os.path.dirname(os.path.realpath(__file__)))
# Load the trained model test
if opts.loadModel != 'none':
model_path = os.path.join(opts.root_dir, opts.loadModel)
model = torch.load(model_path).cuda(device=opts.device)
model.eval()
else:
print('ERROR: No model is loaded!')
return
# Read the input image, pass input to gpu
if opts.img == 'None':
val_dataset = PENN_CROP(opts, 'val')
val_loader = tud.DataLoader(val_dataset,
batch_size=1,
shuffle=False,
num_workers=int(opts.num_workers))
opts.nJoints = val_dataset.nJoints
opts.skeleton = val_dataset.skeleton
for i, gt in enumerate(val_loader):
# Test Visualizer, Input and get_preds
if i == 0:
input, label = gt['input'], gt['label']
gtpts, center, scale, proj = gt['gtpts'], gt['center'], gt[
'scale'], gt['proj']
input_var = input[:, 0, ].float().cuda(device=opts.device,
non_blocking=True)
# output = label
output = model(input_var)
# Test Loss, Err and Acc(PCK)
Loss, Err, Acc = AverageMeter(), AverageMeter(), AverageMeter()
ref = get_ref(opts.dataset, scale)
for j in range(opts.preSeqLen):
pred = get_preds(output[:, j, ].cpu().float())
pred = original_coordinate(pred, center[:, ], scale,
opts.outputRes)
err, ne = error(pred, gtpts[:, j, ], ref)
acc, na = accuracy(pred, gtpts[:, j, ], ref)
# assert ne == na, "ne must be the same as na"
Err.update(err)
Acc.update(acc)
print(j, f"{Err.val:.6f}", Acc.val)
print('all', f"{Err.avg:.6f}", Acc.avg)
# Visualizer Object
## Initialize
v = Visualizer(opts.nJoints, opts.skeleton, opts.outputRes)
# ## Add input image
# v.add_img(input[0,0,].transpose(2, 0).numpy().astype(np.uint8))
# ## Get the predicted joints
# predJoints = get_preds(output[:, 0, ])
# # ## Add joints and skeleton to the figure
# v.add_2d_joints_skeleton(predJoints, (0, 0, 255))
# Transform heatmap to show
hm_img = output[0, 0, ].cpu().detach().numpy()
v.add_hm(hm_img)
## Show image
v.show_img(pause=True)
break
else:
print('NOT ready for the raw input outside the dataset')
img = cv2.imread(opts.img)
input = torch.from_numpy(img.tramspose(2, 0, 1)).float() / 256.
input = input.view(1, input.size(0), input.size(1), input.size(2))
input_var = torch.autograd.variable(input).float().cuda(
device=opts.device)
output = model(input_var)
predJoints = get_preds(output[-2].data.cpu().numpy())[0] * 4