def __getitem__(self, index):
img = self.LoadImage(index)
pts, c, s = self.GetPartInfo(index)
r = 0
if self.split == 'train':
s = s * (2**Rnd(ref.scale))
r = 0 if np.random.random() < 0.6 else Rnd(ref.rotate)
inp = Crop(img, c, s, r, ref.inputRes) / 256.
out = np.zeros((ref.nJoints, ref.outputRes, ref.outputRes))
for i in range(ref.nJoints):
if pts[i][0] > 1:
pt = Transform(pts[i], c, s, r, ref.outputRes)
out[i] = DrawGaussian(out[i], pt, ref.hmGauss)
if self.split == 'train':
if np.random.random() < 0.5:
inp = Flip(inp)
out = ShuffleLR(Flip(out))
inp[0] = np.clip(inp[0] * (np.random.random() * (0.4) + 0.6), 0, 1)
inp[1] = np.clip(inp[1] * (np.random.random() * (0.4) + 0.6), 0, 1)
inp[2] = np.clip(inp[2] * (np.random.random() * (0.4) + 0.6), 0, 1)
meta = np.zeros(1)
else:
meta = {'index': index, 'center': c, 'scale': s, 'rotate': r}
return inp, out, meta
def __getitem__(self, index):
img = self.LoadImage(index)
pts, c, s = self.GetPartInfo(index)
r = 0
if self.split == 'train':
s = s * (2 ** Rnd(ref.scale))
r = 0 if np.random.random() < 0.6 else Rnd(ref.rotate)
inp = Crop(img, c, s, r, ref.inputRes) / 256.
out = np.zeros((ref.nJoints, ref.outputRes, ref.outputRes))
Reg = np.zeros((ref.nJoints, 3))
for i in range(ref.nJoints):
if pts[i][0] > 1:
pt = Transform(pts[i], c, s, r, ref.outputRes)
out[i] = DrawGaussian(out[i], pt, ref.hmGauss)
Reg[i, :2] = pt
Reg[i, 2] = 1
if self.split == 'train':
if np.random.random() < 0.5:
inp = Flip(inp)
out = ShuffleLR(Flip(out))
Reg[:, 1] = Reg[:, 1] * -1
Reg = ShuffleLR(Reg)
#print 'before', inp[0].max(), inp[0].mean()
inp[0] = np.clip(inp[0] * (np.random.random() * (0.4) + 0.6), 0, 1)
inp[1] = np.clip(inp[1] * (np.random.random() * (0.4) + 0.6), 0, 1)
inp[2] = np.clip(inp[2] * (np.random.random() * (0.4) + 0.6), 0, 1)
#print 'after', inp[0].max(), inp[0].mean()
inp = torch.from_numpy(inp)
if self.returnMeta:
return inp, out, Reg, np.zeros((ref.nJoints, 3))
else:
return inp, out
def getitem(self, index, meta, whichFrames, frameWiseData):
img = cv2.imread(ref.posetrackDataDir + "/" + whichFrames[index])
#print(ref.posetrackDataDir + whichFrames[index])
frameData = frameWiseData[whichFrames[index]]
bboxmean, bboxdelta, joints = frameData[meta['personIndex']]
pts = joints
c = np.asarray(bboxmean)
s = bboxdelta * meta['s'] / 100.0
r = meta['r']
inp = Crop(img, c, s, r, ref.inputRes) / 256.
out = np.zeros((ref.nJoints, ref.outputRes, ref.outputRes))
Reg = np.zeros((ref.nJoints, 3))
for i in range(ref.nJoints):
if pts[i][0] > 1:
pt = Transform(pts[i], c, s, r, ref.outputRes)
out[i] = DrawGaussian(out[i], pt, ref.hmGauss)
Reg[i, :2] = pt
Reg[i, 2] = 1
if self.split == 'train':
if meta['flip']:
inp = Flip(inp)
out = ShuffleLR(Flip(out))
Reg[:, 1] = Reg[:, 1] * -1
Reg = ShuffleLR(Reg)
inp[0] = np.clip(inp[0] * (np.random.random() * (0.4) + 0.6), 0, 1)
inp[1] = np.clip(inp[1] * (np.random.random() * (0.4) + 0.6), 0, 1)
inp[2] = np.clip(inp[2] * (np.random.random() * (0.4) + 0.6), 0, 1)
whatever = (np.zeros((ref.nJoints, 3)))
return inp, out, Reg, whatever
def __getitem__(self, index):
img = self.LoadImage(index)
pts, action, c, s = self.GetPartInfo(index)
nb_pts = len(pts)
r = 0
flip = False
if self.split == 'train':
s = s * (2**Rnd(ref.scale))
r = 0 if np.random.random() < 0.6 else Rnd(ref.rotate)
inp = old_div(Crop(img, c, s, r, ref.inputRes), 256.)
if self.split == 'train':
if np.random.random() < 0.5:
inp = Flip(inp)
flip = True
inp[0] = np.clip(inp[0] * (np.random.random() * (0.4) + 0.6), 0, 1)
inp[1] = np.clip(inp[1] * (np.random.random() * (0.4) + 0.6), 0, 1)
inp[2] = np.clip(inp[2] * (np.random.random() * (0.4) + 0.6), 0, 1)
meta = np.zeros(1)
else:
meta = {'index': index, 'center': c, 'scale': s, 'rotate': r}
output = []
for k in range(nb_pts):
out = np.zeros((ref.nJoints, ref.outputRes, ref.outputRes))
for i in range(ref.nJoints):
if pts[k][i][0] > 1:
pt = Transform(pts[k][i], c, s, r, ref.outputRes)
out[i] = DrawGaussian(out[i], pt, ref.hmGauss)
if self.split == 'train':
out = ShuffleLR(Flip(out))
output.append(out)
return inp, output, action, meta
def __getitem__(self, index):
img = self.LoadImage(index)
pts2d, pts3d, emb, c, s = self.GetPartInfo(index)
s = min(s, max(img.shape[0], img.shape[1])) * 1.0
pts3d[:, 2] += s / 2
r = 0
if self.split == 'train':
s = s * (2**Rnd(ref.scale))
c[1] = c[1] + Rnd(ref.shiftY)
r = 0 if np.random.random() < 0.6 else Rnd(ref.rotate)
inp = Crop(img, c, s, r, ref.inputRes)
inp = inp.transpose(2, 0, 1).astype(np.float32) / 256.
starMap = np.zeros((1, ref.outputRes, ref.outputRes))
embMap = np.zeros((3, ref.outputRes, ref.outputRes))
depMap = np.zeros((1, ref.outputRes, ref.outputRes))
mask = np.concatenate([
np.ones((1, ref.outputRes, ref.outputRes)),
np.zeros((4, ref.outputRes, ref.outputRes))
])
for i in range(pts3d.shape[0]):
if self.annot['valid'][index][i] > ref.eps:
if (self.annot['vis'][index][i] > ref.eps):
pt3d = Transform3D(pts3d[i], c, s, r,
ref.outputRes).astype(np.int32)
pt2d = Transform(pts2d[i], c, s, r,
ref.outputRes).astype(np.int32)
if pt2d[0] >= 0 and pt2d[0] < ref.outputRes and pt2d[
1] >= 0 and pt2d[1] < ref.outputRes:
embMap[:, pt2d[1], pt2d[0]] = emb[i]
depMap[0, pt2d[1],
pt2d[0]] = 1.0 * pt3d[2] / ref.outputRes - 0.5
mask[1:, pt2d[1], pt2d[0]] = 1
starMap[0] = np.maximum(
starMap[0],
DrawGaussian(np.zeros((ref.outputRes, ref.outputRes)),
pt2d, ref.hmGauss).copy())
out = starMap
if 'emb' in self.opt.task:
out = np.concatenate([out, embMap])
if 'dep' in self.opt.task:
out = np.concatenate([out, depMap])
mask = mask[:out.shape[0]].copy()
if self.split == 'train':
if np.random.random() < 0.5:
inp = Flip(inp)
out = Flip(out)
mask = Flip(mask)
if 'emb' in self.opt.task:
out[1] = -out[1]
return inp, out, mask
def __getitem__(self, index):
img = self.LoadImage(index)
class_id = self.annot['class_id'][index]
c, s, v = self.GetPartInfo(index)
s = min(s, max(img.shape[0], img.shape[1])) * 1.0
r = 0
if self.split == 'train':
s = s * (2 ** Rnd(ref.scale))
c[1] = c[1] + Rnd(ref.shiftY)
r = 0 if np.random.random() < 0.6 else Rnd(ref.rotate)
v[2] += r / 180.
v[2] += 2 if v[2] < -1 else (-2 if v[2] > 1 else 0)
inp = Crop(img, c, s, r, ref.inputRes)
inp = inp.transpose(2, 0, 1).astype(np.float32) / 256.
if self.split == 'train':
if np.random.random() < 0.5:
inp = Flip(inp)
v[0] = - v[0]
v[2] = - v[2]
v[2] += 2 if v[2] <= -1 else 0
#https://github.com/shubhtuls/ViewpointsAndKeypoints/blob/master/rcnnVp/rcnnBinnedJointTrainValTestCreate.m#L77
vv = v.copy()
if vv[0] < 0:
v[0] = self.opt.numBins - 1 - np.floor(-vv[0] * self.opt.numBins / 2.)
else:
v[0] = np.floor(vv[0] * self.opt.numBins / 2.)
v[1] = np.ceil(vv[1] * self.opt.numBins / 2. + self.opt.numBins / 2. - 1)
v[2] = np.ceil(vv[2] * self.opt.numBins / 2. + self.opt.numBins / 2. - 1)
v = v.astype(np.int32)
if self.opt.specificView:
vv = np.ones(3 * len(ref.pascalClassId), dtype = np.int32) * self.opt.numBins
vv[class_id * 3: class_id * 3 + 3] = v.copy()
v = vv.copy()
return inp, v
def step(split, epoch, opt, dataLoader, model, criterion, optimizer=None):
if split == 'train':
model.train()
else:
model.eval()
Loss, Acc = AverageMeter(), AverageMeter()
preds = []
nIters = len(dataLoader)
bar = Bar('{}'.format(opt.expID), max=nIters)
for i, (input, target, target2, meta) in enumerate(dataLoader):
input_var = torch.autograd.Variable(input).float().cuda()
target_var = torch.autograd.Variable(target).float().cuda()
target_var2 = torch.autograd.Variable(target2).float().cuda()
#print( input_var)
output = model(input_var)
#print(output[-1].size())
if opt.DEBUG >= 2:
gt = getPreds(target.cpu().numpy()) * 4
pred = getPreds((output[opt.nStack - 1].data).cpu().numpy()) * 4
debugger = Debugger()
img = (input[0].numpy().transpose(1, 2, 0) * 256).astype(
np.uint8).copy()
debugger.addImg(img)
debugger.addPoint2D(pred[0], (255, 0, 0))
debugger.addPoint2D(gt[0], (0, 0, 255))
debugger.showAllImg(pause=True)
loss = criterion(output[0], target_var)
for k in range(1, opt.nStack):
loss += criterion(output[k], target_var)
Loss.update(loss.data[0], input.size(0))
Acc.update(
Accuracy((output[opt.nStack - 1].data).cpu().numpy(),
(target_var.data).cpu().numpy()))
if split == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
else:
input_ = input.cpu().numpy()
input_[0] = Flip(input_[0]).copy()
inputFlip_var = torch.autograd.Variable(
torch.from_numpy(input_).view(1, input_.shape[1], ref.inputRes,
ref.inputRes)).float().cuda(
opt.GPU)
outputFlip = model(inputFlip_var)
outputFlip = ShuffleLR(
Flip((outputFlip[opt.nStack -
1].data).cpu().numpy()[0])).reshape(
1, ref.nJoints, 64, 64)
output_ = (
(output[opt.nStack - 1].data).cpu().numpy() + outputFlip) / 2
preds.append(
finalPreds(output_, meta['center'], meta['scale'],
meta['rotate'])[0])
Bar.suffix = '{split} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:} | Loss {loss.avg:.6f} | Acc {Acc.avg:.6f} ({Acc.val:.6f})'.format(
epoch,
i,
nIters,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=Loss,
Acc=Acc,
split=split)
bar.next()
bar.finish()
return {'Loss': Loss.avg, 'Acc': Acc.avg}, preds
def step(split, epoch, opt, dataLoader, model, criterion, optimizer=None):
if split == 'train':
model.train()
else:
model.eval()
preds = []
Loss, Acc = AverageMeter(), AverageMeter()
nIters = len(dataLoader)
bar = Bar('{}'.format(opt.expID), max=nIters)
for i, (input, view) in enumerate(dataLoader):
input_var = torch.autograd.Variable(input.cuda(
opt.GPU, async=True)).float().cuda(opt.GPU)
target_var = torch.autograd.Variable(view.view(-1)).long().cuda(
opt.GPU)
output = model(input_var)
numBins = opt.numBins
loss = torch.nn.CrossEntropyLoss(ignore_index=numBins).cuda(opt.GPU)(
output.view(-1, numBins), target_var)
Acc.update(AccViewCls(output.data, view, numBins, opt.specificView))
Loss.update(loss.data[0], input.size(0))
if split == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
else:
if opt.test:
out = {}
input_ = input.cpu().numpy()
input_[0] = Flip(input_[0]).copy()
inputFlip_var = torch.autograd.Variable(
torch.from_numpy(input_).view(
1, input_.shape[1], ref.inputRes,
ref.inputRes)).float().cuda(opt.GPU)
outputFlip = model(inputFlip_var)
pred = outputFlip.data.cpu().numpy()
numBins = opt.numBins
if opt.specificView:
nCat = len(ref.pascalClassId)
pred = pred.reshape(1, nCat, 3 * numBins)
azimuth = pred[0, :, :numBins]
elevation = pred[0, :, numBins:numBins * 2]
rotate = pred[0, :, numBins * 2:numBins * 3]
azimuth = azimuth[:, ::-1]
rotate = rotate[:, ::-1]
output_flip = []
for c in range(nCat):
output_flip.append(
np.array([azimuth[c], elevation[c],
rotate[c]]).reshape(1, numBins * 3))
output_flip = np.array(output_flip).reshape(
1, nCat * 3 * numBins)
else:
azimuth = pred[0][:numBins]
elevation = pred[0][numBins:numBins * 2]
rotate = pred[0][numBins * 2:numBins * 3]
azimuth = azimuth[::-1]
rotate = rotate[::-1]
output_flip = np.array([azimuth, elevation,
rotate]).reshape(1, numBins * 3)
out['reg'] = (output.data.cpu().numpy() + output_flip) / 2.
preds.append(out)
Bar.suffix = '{split:5} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:} | Loss {loss.avg:.6f} | Acc {Acc.avg:.6f}'.format(
epoch,
i,
nIters,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=Loss,
Acc=Acc,
split=split)
bar.next()
bar.finish()
return {'Loss': Loss.avg, 'Acc': Acc.avg}, preds
def __getitem__(self, index):
seqIdx = self.getSeq(index)
"""
input: predSeqLen x 3 x inputRes x inputRes Input image After Crop and transform
hmap: predSeqLen x numJoints x outputRes x outputRes
gtpts: predSeqLen x numJoints x 2 Joints Positions BEFORE crop and transform
proj: predSeqLen x numJoints x 2 Joints Positions AFTER crop and transform
"""
input = np.zeros((self.nPhase, 3, self.inputRes, self.inputRes))
hmap = np.zeros(
(self.nPhase, self.nJoints, self.outputRes, self.outputRes))
gtpts = np.zeros((self.nPhase, self.nJoints, 2))
repos, trans, focal, proj = {}, {}, {}, {}
for i in range(len(seqIdx)):
sid = seqIdx[i]
im = self.LoadImage(int(sid))
if i == 0:
center, scale = self.getCenterScale(im)
inp = Crop(im, center, scale, 0, self.inputRes)
pts = self.part[int(sid)]
pj = np.zeros(np.shape(pts))
for j in range(len(pts)):
if pts[j][0] != 0 and pts[j][1] != 0:
pj[j] = Transform(pts[j], center, scale, 0, self.outputRes,
False)
hm = np.zeros((np.shape(pts)[0], self.outputRes, self.outputRes))
for j in range(len(pts)):
if pts[j][0] != 0 and pts[j][1] != 0:
DrawGaussian(hm[j], np.round(pj[j]), 2)
inp = inp.transpose(2, 1, 0)
input[i] = inp
repos[i] = np.zeros((np.size(1), 3))
trans[i] = np.zeros(3)
focal[i] = np.zeros(1)
hmap[i] = hm
proj[i] = pj
gtpts[i] = pts
if self.split == 'train':
m1 = np.random.uniform(0.8, 1.2)
m2 = np.random.uniform(0.8, 1.2)
m3 = np.random.uniform(0.8, 1.2)
for i in range(len(input)):
input[i][:, :, 0] = input[i][:, :, 0] * m1
np.clip(input[i][:, :, 0], 0, 1, out=input[i][:, :, 0])
input[i][:, :, 1] = input[i][:, :, 1] * m2
np.clip(input[i][:, :, 1], 0, 1, out=input[i][:, :, 1])
input[i][:, :, 2] = input[i][:, :, 2] * m3
np.clip(input[i][:, :, 2], 0, 1, out=input[i][:, :, 2])
if np.random.uniform() <= 0.5:
for i in range(len(input)):
input[i] = cv2.flip(input[i], 1)
hmap[i] = Flip(ShuffleLR(hmap[i]))
proj[i] = ShuffleLR(proj[i])
ind = np.where(proj[i] == 0)
proj[i][:, 0] = self.outputRes - proj[i][:, 0] + 1
if len(ind[0]) != 0:
proj[i][ind[0][0]] = 0
return {
'input': input,
'label': hmap,
'gtpts': gtpts,
'center': center,
'scale': scale,
'proj': proj
}
def step(split, epoch, opt, dataLoader, model, criterion, optimizer=None):
if split == 'train':
model.train()
else:
model.eval()
Loss, Acc = AverageMeter(), AverageMeter()
preds = []
nIters = len(dataLoader)
bar = Bar('{}'.format(opt.expID), max=nIters)
for i, (input, targets, action, meta) in enumerate(dataLoader):
input_var = torch.autograd.Variable(input).float().cuda(opt.GPU)
target_var = []
for t in range(len(targets)):
target_var.append(
torch.autograd.Variable(targets[t]).float().cuda(opt.GPU))
z = []
for k in range(opt.numNoise):
noise = torch.autograd.Variable(
torch.randn((input_var.shape[0], 1, 64, 64))).cuda(opt.GPU)
z.append(noise)
output, samples = model(input_var, z, action)
pred_sample = maximumExpectedUtility(samples, criterion)
target = maximumExpectedUtility(target_var, criterion)
if opt.DEBUG >= 2:
gt = getPreds(target.cpu().numpy()) * 4
pred = getPreds((pred_sample.data).cpu().numpy()) * 4
debugger = Debugger()
img = (input[0].numpy().transpose(1, 2, 0) * 256).astype(
np.uint8).copy()
debugger.addImg(img)
debugger.addPoint2D(pred[0], (255, 0, 0))
debugger.addPoint2D(gt[0], (0, 0, 255))
debugger.showAllImg(pause=True)
loss = DiscoLoss(output, samples, target_var, criterion)
Loss.update(loss.item(), input.size(0))
Acc.update(
Accuracy((pred_sample.data).cpu().numpy(),
(target.data).cpu().numpy()))
if split == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
else:
input_ = input.cpu().numpy()
input_[0] = Flip(input_[0]).copy()
inputFlip_var = torch.autograd.Variable(
torch.from_numpy(input_).view(1, input_.shape[1], ref.inputRes,
ref.inputRes)).float().cuda(
opt.GPU)
_, samplesFlip = model(inputFlip_var, z, action)
pred_sample_flip = maximumExpectedUtility(samplesFlip, criterion)
outputFlip = ShuffleLR(
Flip((pred_sample_flip.data).cpu().numpy()[0])).reshape(
1, ref.nJoints, ref.outputRes, ref.outputRes)
output_ = old_div(((pred_sample.data).cpu().numpy() + outputFlip),
2)
preds.append(
finalPreds(output_, meta['center'], meta['scale'],
meta['rotate'])[0])
Bar.suffix = '{split} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:} | Loss {loss.avg:.6f} | Acc {Acc.avg:.6f} ({Acc.val:.6f})'.format(
epoch,
i,
nIters,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=Loss,
Acc=Acc,
split=split)
bar.next()
bar.finish()
return {'Loss': Loss.avg, 'Acc': Acc.avg}, preds
def step(split, epoch, opt, dataLoader, model, criterion, optimizer=None):
if split == 'train':
model.train()
else:
model.eval()
preds = []
Loss, LossStar = AverageMeter(), AverageMeter()
nIters = len(dataLoader)
bar = Bar('{}'.format(opt.expID), max=nIters)
for i, (input, target, mask) in enumerate(dataLoader):
if mask.size(1) > 1:
mask[:, 1:, :, :] *= ref.outputRes * (opt.regWeight**0.5)
if opt.GPU > -1:
input_var = torch.autograd.Variable(input.cuda(
opt.GPU, async=True)).float().cuda(opt.GPU)
target_var = torch.autograd.Variable(
target.cuda(opt.GPU, async=True)).float().cuda(opt.GPU)
mask_var = torch.autograd.Variable(mask.cuda(
opt.GPU, async=True)).float().cuda(opt.GPU)
else:
input_var = torch.autograd.Variable(input).float()
target_var = torch.autograd.Variable(target).float()
mask_var = torch.autograd.Variable(mask).float()
output = model(input_var)
output_pred = output[opt.nStack - 1].data.cpu().numpy().copy()
for k in range(opt.nStack):
output[k] = mask_var * output[k]
target_var = mask_var * target_var
loss = 0
for k in range(opt.nStack):
loss += criterion(output[k], target_var)
LossStar.update((
(target.float()[:, 0, :, :] -
output[opt.nStack - 1].cpu().data.float()[:, 0, :, :])**2).mean())
Loss.update(loss.data[0], input.size(0))
if split == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
else:
if opt.test:
out = {}
input_ = input.cpu().numpy()
input_[0] = Flip(input_[0]).copy()
inputFlip_var = torch.autograd.Variable(
torch.from_numpy(input_).view(
1, input_.shape[1], ref.inputRes,
ref.inputRes)).float().cuda(opt.GPU)
outputFlip = model(inputFlip_var)
output_flip = outputFlip[opt.nStack - 1].data.cpu().numpy()
output_flip[0] = Flip(output_flip[0])
if not (opt.task == 'star'):
output_flip[0, 1, :, :] = -output_flip[0, 1, :, :]
output_pred = (output_pred + output_flip) / 2.0
out['map'] = output_pred
preds.append(out)
Bar.suffix = '{split:5} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:} | Loss {loss.avg:.6f} | LossStar {lossStar.avg:.6f}'.format(
epoch,
i,
nIters,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=Loss,
lossStar=LossStar,
split=split)
bar.next()
if opt.DEBUG > 1 or (opt.DEBUG == 1 and i % (nIters / 200) == 0):
for j in range(input.size(0)):
debugger = Debugger()
img = (input[j].numpy()[:3].transpose(1, 2, 0) * 256).astype(
np.uint8).copy()
img2 = img.copy().astype(np.float32)
img3 = img.copy().astype(np.float32)
imgMNS = img.copy()
out = (cv2.resize(
((output[opt.nStack - 1][j, 0].data).cpu().numpy()).copy(),
(ref.inputRes, ref.inputRes)) * 256)
gtmap = (cv2.resize((target[j, 0].cpu().numpy()).copy(),
(ref.inputRes, ref.inputRes)) * 256)
out[out < 0] = 0
out[out > 255] = 255
img2[:, :, 0] = (img2[:, :, 0] + out)
img2[img2 > 255] = 255
img3[:, :, 2] = (img3[:, :, 2] + gtmap)
img3[img3 > 255] = 255
gtmap[gtmap > 255] = 255
idx = i * input.size(0) + j if opt.DEBUG == 1 else 0
img2, out, gtmap, img3 = img2.astype(np.uint8), out.astype(
np.uint8), gtmap.astype(np.uint8), img3.astype(np.uint8)
if 'emb' in opt.task:
gt, pred = [], []
ps = parseHeatmap(target[j].numpy())
print('ps', ps)
for k in range(len(ps[0])):
print('target', k, target[j, 1:4, ps[0][k],
ps[1][k]].numpy())
x, y, z = (
(target[j, 1:4, ps[0][k], ps[1][k]].numpy() + 0.5)
* 255).astype(np.int32)
gt.append(target[j, 1:4, ps[0][k], ps[1][k]].numpy())
cv2.circle(imgMNS, (ps[1][k] * 4, ps[0][k] * 4), 6,
(int(x), int(y), int(z)), -1)
ps = parseHeatmap(output_pred[j])
for k in range(len(ps[0])):
print('pred', k, output_pred[j, 1:4, ps[0][k],
ps[1][k]])
x, y, z = (
(output_pred[j, 1:4, ps[0][k], ps[1][k]] + 0.5) *
255).astype(np.int32)
pred.append(output_pred[j, 1:4, ps[0][k], ps[1][k]])
cv2.circle(imgMNS, (ps[1][k] * 4, ps[0][k] * 4), 4,
(255, 255, 255), -1)
cv2.circle(imgMNS, (ps[1][k] * 4, ps[0][k] * 4), 2,
(int(x), int(y), int(z)), -1)
debugger.addPoint3D(np.array(gt), c='auto', marker='o')
#debugger.addPoint3D(np.array(pred), c = 'auto', marker = 'x')
debugger.addImg(imgMNS, '{}_mns'.format(idx))
debugger.addImg(out, '{}_out'.format(idx))
debugger.addImg(gtmap, '{}_gt'.format(idx))
debugger.addImg(img, '{}_img'.format(idx))
debugger.addImg(img2, '{}_img2'.format(idx))
debugger.addImg(img3, '{}_img3'.format(idx))
if opt.DEBUG == 1:
debugger.saveAllImg(path=opt.debugPath)
else:
debugger.showAllImg(pause=not ('emb' in opt.task))
if 'emb' in opt.task:
debugger.show3D()
bar.finish()
return {'Loss': Loss.avg, 'LossStar': LossStar.avg}, preds