class ConvPoseMachine(object): """ The class for Pose Estimation. Include train, val, val & predict. """ def __init__(self, configer): self.configer = configer self.batch_time = AverageMeter() self.data_time = AverageMeter() self.train_losses = AverageMeter() self.val_losses = AverageMeter() self.pose_visualizer = PoseVisualizer(configer) self.pose_model_manager = ModelManager(configer) self.pose_data_loader = DataLoader(configer) self.pose_net = None self.train_loader = None self.val_loader = None self.optimizer = None self.scheduler = None self.runner_state = dict() self._init_model() def _init_model(self): self.pose_net = self.pose_model_manager.get_single_pose_model() self.pose_net = RunnerHelper.load_net(self, self.pose_net) self.optimizer, self.scheduler = Trainer.init( self._get_parameters(), self.configer.get('solver')) self.train_loader = self.pose_data_loader.get_trainloader() self.val_loader = self.pose_data_loader.get_valloader() self.cpm_loss = self.pose_model_manager.get_pose_loss() def _get_parameters(self): return self.pose_net.parameters() def train(self): """ Train function of every epoch during train phase. """ self.pose_net.train() start_time = time.time() # Adjust the learning rate after every epoch. self.runner_state['epoch'] += 1 # data_tuple: (inputs, heatmap, maskmap, tagmap, num_objects) for i, data_dict in enumerate(self.train_loader): Trainer.update(self, solver_dict=self.configer.get('solver')) self.data_time.update(time.time() - start_time) # Change the data type. # Forward pass. out_dict = self.pose_net(data_dict) # Compute the loss of the train batch & backward. loss = self.cpm_loss(out_dict, data_dict, gathered=self.configer.get( 'network', 'gathered')) self.train_losses.update(loss.item(), len(DCHelper.tolist(data_dict['meta']))) self.optimizer.zero_grad() loss.backward() self.optimizer.step() # Update the vars of the train phase. self.batch_time.update(time.time() - start_time) start_time = time.time() self.runner_state['iters'] += 1 # Print the log info & reset the states. if self.runner_state['iters'] % self.configer.get( 'solver', 'display_iter') == 0: Log.info( 'Train Epoch: {0}\tTrain Iteration: {1}\t' 'Time {batch_time.sum:.3f}s / {2}iters, ({batch_time.avg:.3f})\t' 'Data load {data_time.sum:.3f}s / {2}iters, ({data_time.avg:3f})\n' 'Learning rate = {3}\tLoss = {loss.val:.8f} (ave = {loss.avg:.8f})\n' .format(self.runner_state['epoch'], self.runner_state['iters'], self.configer.get('solver', 'display_iter'), RunnerHelper.get_lr(self.optimizer), batch_time=self.batch_time, data_time=self.data_time, loss=self.train_losses)) self.batch_time.reset() self.data_time.reset() self.train_losses.reset() if self.configer.get('solver', 'lr')['metric'] == 'iters' \ and self.runner_state['iters'] == self.configer.get('solver', 'max_iters'): break # Check to val the current model. if self.runner_state['iters'] % self.configer.get( 'solver', 'test_interval') == 0: self.val() def val(self): """ Validation function during the train phase. """ self.pose_net.eval() start_time = time.time() with torch.no_grad(): for j, data_dict in enumerate(self.val_loader): # Forward pass. out_dict = self.pose_net(data_dict) # Compute the loss of the val batch. loss = self.cpm_loss(out_dict, data_dict, gathered=self.configer.get( 'network', 'gathered')) self.val_losses.update(loss.item(), len(DCHelper.tolist(data_dict['meta']))) # Update the vars of the val phase. self.batch_time.update(time.time() - start_time) start_time = time.time() RunnerHelper.save_net(self, self.pose_net, iters=self.runner_state['iters']) # Print the log info & reset the states. Log.info( 'Test Time {batch_time.sum:.3f}s, ({batch_time.avg:.3f})\t' 'Loss {loss.avg:.8f}\n'.format(batch_time=self.batch_time, loss=self.val_losses)) self.batch_time.reset() self.val_losses.reset() self.pose_net.train()
class ConvPoseMachineTest(object): def __init__(self, configer): self.configer = configer self.blob_helper = BlobHelper(configer) self.pose_vis = PoseVisualizer(configer) self.pose_model_manager = ModelManager(configer) self.pose_data_loader = DataLoader(configer) self.device = torch.device('cpu' if self.configer.get('gpu') is None else 'cuda') self.pose_net = None self._init_model() def _init_model(self): self.pose_net = self.pose_model_manager.get_single_pose_model() self.pose_net = RunnerHelper.load_net(self, self.pose_net) self.pose_net.eval() def __test_img(self, image_path, save_path): Log.info('Image Path: {}'.format(image_path)) ori_image = ImageHelper.read_image(image_path, tool=self.configer.get('data', 'image_tool'), mode=self.configer.get('data', 'input_mode')) ori_width, ori_height = ImageHelper.get_size(ori_image) ori_img_bgr = ImageHelper.get_cv2_bgr(ori_image, mode=self.configer.get('data', 'input_mode')) heatmap_avg = np.zeros((ori_height, ori_width, self.configer.get('network', 'heatmap_out'))) for i, scale in enumerate(self.configer.get('test', 'scale_search')): image = self.blob_helper.make_input(ori_image, input_size=self.configer.get('test', 'input_size'), scale=scale) with torch.no_grad(): heatmap_out_list = self.pose_net(image) heatmap_out = heatmap_out_list[-1] # extract outputs, resize, and remove padding heatmap = heatmap_out.squeeze(0).cpu().numpy().transpose(1, 2, 0) heatmap = cv2.resize(heatmap, (ori_width, ori_height), interpolation=cv2.INTER_CUBIC) heatmap_avg = heatmap_avg + heatmap / len(self.configer.get('test', 'scale_search')) all_peaks = self.__extract_heatmap_info(heatmap_avg) image_canvas = self.__draw_key_point(all_peaks, ori_img_bgr) ImageHelper.save(image_canvas, save_path) def __extract_heatmap_info(self, heatmap_avg): all_peaks = [] for part in range(self.configer.get('network', 'heatmap_out') - 1): map_ori = heatmap_avg[:, :, part] map_gau = gaussian_filter(map_ori, sigma=3) map_left = np.zeros(map_gau.shape) map_left[1:, :] = map_gau[:-1, :] map_right = np.zeros(map_gau.shape) map_right[:-1, :] = map_gau[1:, :] map_up = np.zeros(map_gau.shape) map_up[:, 1:] = map_gau[:, :-1] map_down = np.zeros(map_gau.shape) map_down[:, :-1] = map_gau[:, 1:] peaks_binary = np.logical_and.reduce( (map_gau >= map_left, map_gau >= map_right, map_gau >= map_up, map_gau >= map_down, map_gau > self.configer.get('vis', 'part_threshold'))) peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) # note reverse peaks = list(peaks) peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks] all_peaks.append(peaks_with_score) return all_peaks def __draw_key_point(self, all_peaks, img_raw): img_canvas = img_raw.copy() # B,G,R order for i in range(self.configer.get('network', 'heatmap_out') - 1): for j in range(len(all_peaks[i])): cv2.circle(img_canvas, all_peaks[i][j][0:2], self.configer.get('vis', 'stick_width'), self.configer.get('details', 'color_list')[i], thickness=-1) return img_canvas def debug(self, vis_dir): for i, data_dict in enumerate(self.pose_data_loader.get_trainloader()): inputs = data_dict['img'] heatmap = data_dict['heatmap'] for j in range(inputs.size(0)): image_bgr = self.blob_helper.tensor2bgr(inputs[j]) heatmap_avg = heatmap[j].numpy().transpose(1, 2, 0) heatmap_avg = cv2.resize(heatmap_avg, (0, 0), fx=self.configer.get('network', 'stride'), fy=self.configer.get('network', 'stride'), interpolation=cv2.INTER_CUBIC) all_peaks = self.__extract_heatmap_info(heatmap_avg) image_save = self.__draw_key_point(all_peaks, image_bgr) cv2.imwrite(os.path.join(vis_dir, '{}_{}_result.jpg'.format(i, j)), image_save)