def __getitem__(self, index):
if self.split == 'train':
index = np.random.randint(self.nSamples)
img = self.LoadImage(index)
pts, c, s, pts_3d, pts_3d_mono = self.GetPartInfo(index)
pts_3d[7] = (pts_3d[12] +
pts_3d[13]) / 2 # neck = average of shoulders
inp = Crop(img, c, s, 0,
ref.inputRes) / 256. # crop image to input resolution
outMap = np.zeros((ref.nJoints, ref.outputRes,
ref.outputRes)) # nJoints x 64 x 64 heatmap for 2D
outReg = np.zeros((ref.nJoints, 3)) # Regression target for 3D
for i in range(ref.nJoints):
pt = Transform3D(pts_3d[i], c, s, 0, ref.outputRes)
if pts[i][0] > 1:
outMap[i] = DrawGaussian(
outMap[i], pt[:2],
ref.hmGauss) # Draw 2D heat map for detection
outReg[i, 2] = pt[2] / ref.outputRes * 2 - 1
inp = torch.from_numpy(inp)
return inp, outMap, outReg, pts_3d_mono
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):
if self.split == 'train':
index = np.random.randint(self.nSamples)
pts, c, s, pts_3d, pts_3d_mono = self.GetPartInfo(index)
pts_3d[7] = (pts_3d[12] + pts_3d[13]) / 2
out3d = np.zeros((ref.nJoints, 3))
for i in range(ref.nJoints):
pt = Transform3D(pts_3d[i], c, s, 0, ref.outputRes)
out3d[i, :2] = pt[:2]
out3d[i, 2] = pt[2]
return out3d, pts_3d_mono
def __getitem__(self, index):
if self.split == 'train':
index = np.random.randint(self.nSamples)
img = self.LoadImage(index)
pts, c, s, pts_3d, pts_3d_mono = self.GetPartInfo(index)
pts_3d[7] = (pts_3d[12] + pts_3d[13]) / 2
inp = Crop(img, c, s, 0, ref.inputRes) / 256.
outMap = np.zeros((ref.nJoints, ref.outputRes, ref.outputRes))
outReg = np.zeros((ref.nJoints, 3))
for i in range(ref.nJoints):
pt = Transform3D(pts_3d[i], c, s, 0, ref.outputRes)
if pts[i][0] > 1:
outMap[i] = DrawGaussian(outMap[i], pt[:2], ref.hmGauss)
outReg[i, 2] = pt[2] / ref.outputRes * 2 - 1
inp = torch.from_numpy(inp)
return inp, outMap, outReg, pts_3d_mono
def LoadFrameAndData(self, path, frameName):
frame = cv2.imread(path + frameName)
pts_2d, pts_3d, pts_3d_mono = pickle.load(open(
path + "data.pkl", 'rb'))[int(frameName[-10:-4])]
pts_2d = pts_2d
pts_3d = pts_3d
pts_3d_mono = pts_3d_mono
c = np.ones(2) * ref.h36mImgSize / 2
s = ref.h36mImgSize * 1.0
pts_3d = pts_3d - pts_3d[self.root]
s2d, s3d = 0, 0
for e in ref.edges:
s2d += ((pts_2d[e[0]] - pts_2d[e[1]])**2).sum()**0.5
s3d += ((pts_3d[e[0], :2] - pts_3d[e[1], :2])**2).sum()**0.5
scale = s2d / s3d
for j in range(ref.nJoints):
pts_3d[j, 0] = pts_3d[j, 0] * scale + pts_2d[self.root, 0]
pts_3d[j, 1] = pts_3d[j, 1] * scale + pts_2d[self.root, 1]
pts_3d[j, 2] = pts_3d[j, 2] * scale + ref.h36mImgSize / 2
pts_3d[7, :] = (pts_3d[12, :] + pts_3d[13, :]) / 2
frame = Crop(frame, c, s, 0, ref.inputRes) / 256.
outMap = np.zeros((ref.nJoints, ref.outputRes, ref.outputRes))
outReg = np.zeros((ref.nJoints, 3))
for i in range(ref.nJoints):
pt = Transform3D(pts_3d[i], c, s, 0, ref.outputRes)
if pts_2d[i][0] > 1:
outMap[i] = DrawGaussian(outMap[i], pt[:2], ref.hmGauss)
outReg[i, 2] = pt[2] / ref.outputRes * 2 - 1
return frame, outMap, pts_2d, outReg, pts_3d_mono
def LoadFrameAndData(self, path, frameName):
frame = cv2.imread(path + frameName)
dict = pickle.load(open(path + "data.pkl", 'rb'))
pts_2d = pts_2d['2d'][int(frameName[:-4]), :]
pts_3d = pts_3d['3d'][int(frameName[:-4]), :]
pts_3d_mono = pts_3d_mono['3d'][int(frameName[:-4]), :]
c = np.ones(2) * ref.ntuImgSize / 2
s = ref.ntuImgSize * 1.0
pts_3d = pts_3d - pts_3d[self.root]
s2d, s3d = 0, 0
for e in ref.edges:
s2d += ((pts_2d[e[0]] - pts_2d[e[1]])**2).sum()**0.5
s3d += ((pts_3d[e[0], :2] - pts_3d[e[1], :2])**2).sum()**0.5
scale = s2d / s3d
for j in range(ref.nJoints):
pts_3d[j, 0] = pts_3d[j, 0] * scale + pts_2d[self.root, 0]
pts_3d[j, 1] = pts_3d[j, 1] * scale + pts_2d[self.root, 1]
pts_3d[j, 2] = pts_3d[j, 2] * scale + ref.h36mImgSize / 2
pts_3d[7, :] = (pts_3d[12, :] + pts_3d[13, :]) / 2
frame = Crop(frame, c, s, 0, ref.inputRes) / 256.
outReg = np.zeros((ref.nJoints, 3))
for i in range(ref.nJoints):
pt = Transform3D(pts_3d[i], c, s, 0, ref.outputRes)
outReg[i, 2] = pt[2] / ref.outputRes * 2 - 1
frame = torch.from_numpy(frame)
pts_2d = torch.from_numpy(pts_2d)
outReg = torch.from_numpy(outReg)
pts_3d_mono = torch.from_numpy(pts_3d_mono)
return frame, pts_2d, outReg, pts_3d_mono