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