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)