def step(split, epoch, opt, dataLoader, model, criterion, optimizer=None):
if split == 'train':
model.train()
else:
model.eval()
Loss, Acc, Mpjpe, Loss3D = AverageMeter(), AverageMeter(), AverageMeter(
), AverageMeter()
nIters = len(dataLoader)
bar = Bar('==>', max=nIters)
for i, (input, target2D, target3D, meta) in enumerate(dataLoader):
input_var = torch.autograd.Variable(input).float().cuda()
target2D_var = torch.autograd.Variable(target2D).float().cuda()
target3D_var = torch.autograd.Variable(target3D).float().cuda()
depMap, depth = model(input_var, target2D_var)
depthPridict = depth[opt.nStack - 1]
if opt.DEBUG >= 2:
gt = getPreds(target2D.cpu().numpy()) * 4
pred = getPreds((depMap[opt.nStack - 1].data).cpu().numpy()) * 4
debugger = Debugger()
debugger.addImg(
(input[0].numpy().transpose(1, 2, 0) * 256).astype(np.uint8))
debugger.addPoint2D(pred[0], (255, 0, 0))
debugger.addPoint2D(gt[0], (0, 0, 255))
debugger.showImg()
debugger.saveImg('debug/{}.png'.format(i))
loss = 0
for k in range(opt.nStack):
loss += criterion(depth[k], target3D_var[:, :, 2])
Loss.update(loss.data[0], input.size(0))
#Acc.update(Accuracy((output[opt.nStack - 1].data).cpu().numpy(), (target2D_var.data).cpu().numpy()))
mpjpe, num3D = MPJPE2(target3D.cpu().numpy(),
(depthPridict.data).cpu().numpy(), meta)
Mpjpe.update(mpjpe, num3D)
if split == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
Bar.suffix = '{split} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:} | Loss {loss.avg:.6f} | Mpjpe {Mpjpe.avg:.6f} ({Mpjpe.val:.6f})'.format(
epoch,
i,
nIters,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=Loss,
split=split,
Mpjpe=Mpjpe)
bar.next()
bar.finish()
return Loss.avg, Acc.avg, Mpjpe.avg, Loss3D.avg
def main():
opt = opts().parse()
if opt.loadModel != 'none':
model = torch.load(opt.loadModel).cuda()
else:
model = torch.load('hgreg-3d.pth').cuda()
img = cv2.imread(opt.demo)
input = torch.from_numpy(img.transpose(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()
output = model(input_var)
pred = getPreds((output[-2].data).cpu().numpy())[0] * 4
reg = (output[-1].data).cpu().numpy().reshape(pred.shape[0], 1)
debugger = Debugger()
debugger.addImg((input[0].numpy().transpose(1, 2, 0)*256).astype(np.uint8))
debugger.addPoint2D(pred, (255, 0, 0))
debugger.addPoint3D(np.concatenate([pred, (reg + 1) / 2. * 256], axis = 1))
debugger.showImg(pause = True)
debugger.show3D()
def generate(imageName):
process_image(imageName)
model = torch.load('../model/Stage3/model_10.pth',
map_location=lambda storage, loc: storage)
img = cv2.imread(imageName)
input = torch.from_numpy(img.transpose(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()
output = model(input_var)
print(output[-2].data[0][-2].shape)
pred = getPreds((output[-2].data).cpu().numpy())[0] * 4
reg = (output[-1].data).cpu().numpy().reshape(pred.shape[0], 1)
print(pred, (reg + 1) / 2. * 256)
debugger = Debugger()
debugger.addImg(
(input[0].numpy().transpose(1, 2, 0) * 256).astype(np.uint8))
debugger.addPoint2D(pred, (255, 0, 0))
debugger.addPoint3D(np.concatenate([pred, (reg + 1) / 2. * 256], axis=1))
debugger.showImg(pause=True)
debugger.show3D()
'''
scale = 1.0 * h / mheight
new_im = image.resize((int(w / scale), int(h / scale)),
Image.ANTIALIAS)
new_im.save(filename)
new_im.close()
#opt = opts().parse()
imageName = './images/test3.jpg'
#process_image(imageName)
model = torch.load('../model/Stage3/model_10.pth',
map_location=lambda storage, loc: storage)
img = cv2.imread(imageName)
print(type(np.array(img)))
input = torch.from_numpy(img.transpose(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()
output = model(input_var)
pred = getPreds((output[-2].data).cpu().numpy())[0] * 4
reg = (output[-1].data).cpu().numpy().reshape(pred.shape[0], 1)
print(pred, (reg + 1) / 2. * 256)
debugger = Debugger()
debugger.addImg((input[0].numpy().transpose(1, 2, 0) * 256).astype(np.uint8))
debugger.addPoint2D(pred, (255, 0, 0))
debugger.addPoint3D(np.concatenate([pred, (reg + 1) / 2. * 256], axis=1))
debugger.showImg(pause=True)
debugger.show3D()
def step(split, epoch, opt, dataLoader, model, criterion, optimizer=None):
if split == 'train':
model.train()
else:
model.eval()
Loss, Acc, Mpjpe, Loss3D = AverageMeter(), AverageMeter(), AverageMeter(
), AverageMeter()
nIters = len(dataLoader)
bar = Bar('==>', max=nIters)
for i, (input, target2D, target3D, meta) in enumerate(dataLoader):
input_var = torch.autograd.Variable(input).float().cuda()
target2D_var = torch.autograd.Variable(target2D).float().cuda()
target3D_var = torch.autograd.Variable(target3D).float().cuda()
output = model(input_var)
reg = output[opt.nStack]
if opt.DEBUG >= 2:
gt = getPreds(target2D.cpu().numpy()) * 4
pred = getPreds((output[opt.nStack - 1].data).cpu().numpy()) * 4
debugger = Debugger()
debugger.addImg(
(input[0].numpy().transpose(1, 2, 0) * 256).astype(np.uint8))
debugger.addPoint2D(pred[0], (255, 0, 0))
debugger.addPoint2D(gt[0], (0, 0, 255))
debugger.showImg()
debugger.saveImg('debug/{}.png'.format(i))
# FusioCriterion is an Autograd funciton which can be called only once in the forward pass. So it is defined again in every iteration.
# Don't ask why.
loss = FusionCriterion(opt.regWeight, opt.varWeight)(reg, target3D_var)
Loss3D.update(loss.data[0], input.size(0))
for k in range(opt.nStack):
loss += criterion(output[k], target2D_var)
Loss.update(loss.data[0], input.size(0))
Acc.update(
Accuracy((output[opt.nStack - 1].data).cpu().numpy(),
(target2D_var.data).cpu().numpy()))
mpjpe, num3D = MPJPE((output[opt.nStack - 1].data).cpu().numpy(),
(reg.data).cpu().numpy(), meta)
if num3D > 0:
Mpjpe.update(mpjpe, num3D)
if split == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
Bar.suffix = '{split} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:} | Loss {loss.avg:.6f} | Loss3D {loss3d.avg:.6f} | Acc {Acc.avg:.6f} | Mpjpe {Mpjpe.avg:.6f} ({Mpjpe.val:.6f})'.format(
epoch,
i,
nIters,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=Loss,
Acc=Acc,
split=split,
Mpjpe=Mpjpe,
loss3d=Loss3D)
bar.next()
bar.finish()
return Loss.avg, Acc.avg, Mpjpe.avg, Loss3D.avg
