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 = AlexNet(ref.nJoints).cuda() model.load_state_dict(torch.load("save.model")) for (i, filename) in enumerate(os.listdir("./testimages/")): img = cv2.imread("./testimages/" + filename) c = np.ones(2) * ref.h36mImgSize / 2 s = ref.h36mImgSize * 1.0 img2 = Crop(img, c, s, 0, ref.inputRes) / 256. input = torch.from_numpy(img2) input = input.contiguous().view(1, input.size(0), input.size(1), input.size(2)) print(input.size()) input_var = torch.autograd.Variable(input).float().cuda() output = model(input_var) print(output.size()) reg = (output.data).cpu().numpy() #.reshape(pred.shape[0], 1) four = lambda t: t * 4.57 fourfunc = np.vectorize(four) reg = fourfunc(reg) print(reg) debugger = Debugger() debugger.addImg( (input[0].numpy().transpose(1, 2, 0) * 256).astype(np.uint8)) debugger.addPoint2D(reg, (255, 0, 0)) #debugger.addPoint3D(np.concatenate([pred, (reg + 1) / 2. * 256], axis = 1)) debugger.saveImg(path="./result/" + filename) np.set_printoptions(threshold=np.inf, linewidth=np.inf) with open("./result/" + filename[:-4] + ".out", 'w') as f: f.write(np.array2string(reg, separator=', ')) """
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() debugger.saveImg() print(pred) print(len(output),output[0].data.cpu().numpy().shape)
def main(): pickle.load = partial(pickle.load, encoding="latin1") pickle.Unpickler = partial(pickle.Unpickler, encoding="latin1") opt = opts().parse() if opt.loadModel != 'none': model = torch.load(opt.loadModel).cuda() else: model = torch.load('../../tr_models/hgreg-3d.pth').cuda() #opt.demo has the path to dir containing frames of demo video all_frames = os.listdir(opt.demo) n_frames = len(all_frames) #specifics dir_name = opt.demo.split('/')[-1] save_path = '../../output/demo/'+dir_name try: os.makedirs(save_path) except OSError: pass for idx, frame in enumerate(all_frames): print('processing frame {}'.format(idx)) img = cv2.imread(opt.demo+'/'+frame) 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.saveImg(path=save_path+'/frame{}.jpg'.format(idx)) debugger.save3D(path=save_path+'/frame_p3d{}.jpg'.format(idx)) print('frame {} done'.format(idx))
class PoseExtractor: def __init__(self, flag_save_pose_image, flag_save_pose_file): rospy.loginfo("Initializing 3D Pose Extractor") self.frameInfo = FrameInfo() self.bridge = CvBridge() self.lock = Lock() self.image_shape = (256, 256, 3) self.debugger = Debugger() self.tracking_info_topic = rospy.get_param( '~tracking_info_topic', '/person_tracker/tracking_info') self.model_name = rospy.get_param('~pose_model', 'hgreg-3d.pth') self.model = {} self.save_pose_image = flag_save_pose_image self.save_pose_file = flag_save_pose_file self.publish_person = False self.initModel() self.frame_info_pub = rospy.Publisher('~frame_info', FrameInfo, queue_size=1) self.person_pub = rospy.Publisher('~person', Person, queue_size=1) self.tracking_info_sub = rospy.Subscriber(self.tracking_info_topic, FrameInfo, self.trackingInfoCallback, queue_size=1) def trackingInfoCallback(self, tracking_info_msg): begin = time.time() self.frameInfo.frame_id = tracking_info_msg.frame_id self.frameInfo.image_frame = tracking_info_msg.image_frame self.frameInfo.last_frame = tracking_info_msg.last_frame rospy.loginfo("Frame ID: {}".format(self.frameInfo.frame_id)) persons = tracking_info_msg.persons numPersons = len(persons) if numPersons != 0: for person in persons: rospy.loginfo("Person {} is detected".format(person.person_id)) try: # multi-threading for publishing single person #p = ThreadPool(numPersons) #p.map(self.poseEstimation, persons) #p.close() for person in persons: self.poseEstimation(person) self.frame_info_pub.publish(self.frameInfo) self.frameInfo = FrameInfo() if tracking_info_msg.last_frame: rospy.loginfo('Last frame in the video!') except BaseException as e: rospy.logerr(e) else: rospy.logwarn("No person is detected!") self.frame_info_pub.publish(self.frameInfo) if tracking_info_msg.last_frame: rospy.loginfo('Last frame in the video!') rospy.loginfo("FPS: {}".format(1 / (time.time() - begin))) 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)) # Initialize model def initModel(self): rospy.loginfo("=====> Loading and Initializing Model") model_path = rospkg.RosPack().get_path( 'pose_3d_ros') + '/models/' + self.model_name self.model = torch.load(model_path).cuda() img = np.zeros((256, 256, 3)) 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 = self.model(input_var) rospy.loginfo("Model Initialization Done")
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