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) 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 t_inp = inp.transpose(1, 2, 0) out = cv2.resize(t_inp, dsize=(64, 64)) out = torch.from_numpy(out.transpose(2, 0, 1)) inp = torch.from_numpy(inp) return inp, out, outMap
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 poseEstimation(self, tracked_person): person_id = tracked_person.person_id try: curImage = self.bridge.imgmsg_to_cv2(self.frameInfo.image_frame) person_image = curImage[ int(tracked_person.bbox.top):int(tracked_person.bbox.top + tracked_person.bbox.height), int(tracked_person.bbox.left):int(tracked_person.bbox.left + tracked_person.bbox.width)] except CvBridgeError as e: rospy.logerr(e) # Resize input image rospy.logdebug("person image shape: {}".format(person_image.shape)) if person_image.shape != self.image_shape: h, w = person_image.shape[0], person_image.shape[1] center = torch.FloatTensor((w / 2, h / 2)) scale = 1.0 * max(h, w) res = 256 input_image = Crop(person_image, center, scale, 0, res) else: input_image = person_image # Feed input image to model rospy.loginfo("feeding image to model") input = torch.from_numpy(input_image.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() # lock when using model to estimate pose self.lock.acquire() try: output = self.model(input_var) finally: self.lock.release() rospy.logdebug("got output from model") # Get 2D pose rospy.logdebug("Rendering 2D pose") pose2D = getPreds((output[-2].data).cpu().numpy())[0] * 4 # Get 3D pose rospy.logdebug("Rendering 3D pose") reg = (output[-1].data).cpu().numpy().reshape(pose2D.shape[0], 1) pose3D = np.concatenate([pose2D, (reg + 1) / 2. * 256], axis=1) rospy.logdebug("pose 3d shape: {}".format(pose3D.shape)) for pose in pose3D: joint = Point() joint.x = pose[0] joint.y = pose[1] joint.z = pose[2] tracked_person.person_pose.append(joint) # publish person if self.publish_person: self.person_pub.publish(tracked_person) self.lock.acquire() try: self.frameInfo.persons.append(tracked_person) finally: self.lock.release() rospy.logdebug("pose3D: \n {}".format(pose3D)) # Save pose image if self.save_pose_image: cv2.imwrite( pkg_path + '/scripts/debug/original/ogImg_' + str(self.frame_id) + '.png', self.cv_image) cv2.imwrite( pkg_path + '/scripts/debug/input/inputImg_' + str(self.frame_id) + '.png', input_image) self.debugger.addImg(input_image, imgId=self.frame_id) self.debugger.addPoint2D(pose2D, (255, 0, 0), imgId=self.frame_id) self.debugger.saveImg(pkg_path + '/scripts/debug/pose/poseImg_' + str(self.frame_id) + '.png', imgId=self.frame_id) if self.save_pose_file: file_name = pkg_path + '/pose_file/pose_{:04d}.txt'.format( self.frame_id) with file(file_name, 'w') as outfile: np.savetxt(outfile, pose3D, fmt='%-7.2f') rospy.loginfo("Person {} processing finished".format(person_id))
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 }