def tensor2bgr(self, tensor): assert len(tensor.size()) == 3 ori_img = DeNormalize(**self.configer.get('data', 'normalize'))( tensor.cpu()) ori_img = ori_img.numpy().transpose(1, 2, 0).astype(np.uint8) if self.configer.get('data', 'input_mode') == 'BGR': return ori_img else: image_bgr = cv2.cvtColor(ori_img, cv2.COLOR_RGB2BGR) return image_bgr
def vis_default_bboxes(self, ori_img_in, default_bboxes, labels, name='default', sub_dir='encode'): base_dir = os.path.join(self.configer.get('project_dir'), DET_DIR, sub_dir) if not os.path.exists(base_dir): log.error('Dir:{} not exists!'.format(base_dir)) os.makedirs(base_dir) if not isinstance(ori_img_in, np.ndarray): ori_img = DeNormalize( div_value=self.configer.get('normalize', 'div_value'), mean=self.configer.get('normalize', 'mean'), std=self.configer.get('normalize', 'std'))(ori_img_in.clone()) ori_img = ori_img.data.cpu().squeeze().numpy().transpose( 1, 2, 0).astype(np.uint8) ori_img = cv2.cvtColor(ori_img, cv2.COLOR_RGB2BGR) else: ori_img = ori_img_in.copy() assert labels.size(0) == default_bboxes.size(0) bboxes = torch.cat([ default_bboxes[:, :2] - default_bboxes[:, 2:] / 2, default_bboxes[:, :2] + default_bboxes[:, 2:] / 2 ], 1) height, width, _ = ori_img.shape for i in range(labels.size(0)): if labels[i] == 0: continue class_name = self.configer.get('details', 'name_seq')[labels[i] - 1] color_num = len(self.configer.get('details', 'color_list')) cv2.rectangle( ori_img, (int(bboxes[i][0] * width), int(bboxes[i][1] * height)), (int(bboxes[i][2] * width), int(bboxes[i][3] * height)), color=self.configer.get( 'details', 'color_list')[(labels[i] - 1) % color_num], thickness=3) cv2.putText(ori_img, class_name, (int(bboxes[i][0] * width) + 5, int(bboxes[i][3] * height) - 5), cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=self.configer.get('details', 'color_list')[(labels[i] - 1) % color_num], thickness=2) img_path = os.path.join(base_dir, '{}.jpg'.format(name)) cv2.imwrite(img_path, ori_img)
def vis_rois(self, inputs, indices_and_rois, rois_labels=None, name='default', sub_dir='rois'): base_dir = os.path.join(self.configer.get('project_dir'), DET_DIR, sub_dir) if not os.path.exists(base_dir): log.error('Dir:{} not exists!'.format(base_dir)) os.makedirs(base_dir) for i in range(inputs.size(0)): rois = indices_and_rois[indices_and_rois[:, 0] == i][:, 1:] ori_img = DeNormalize(div_value=self.configer.get('normalize', 'div_value'), mean=self.configer.get('normalize', 'mean'), std=self.configer.get('normalize', 'std'))(inputs[i]) ori_img = ori_img.data.cpu().squeeze().numpy().transpose(1, 2, 0).astype(np.uint8) ori_img = cv2.cvtColor(ori_img, cv2.COLOR_RGB2BGR) color_num = len(self.configer.get('details', 'color_list')) for j in range(len(rois)): label = 1 if rois_labels is None else rois_labels[j] if label == 0: continue class_name = self.configer.get('details', 'name_seq')[label - 1] cv2.rectangle(ori_img, (int(rois[j][0]), int(rois[j][1])), (int(rois[j][2]), int(rois[j][3])), color=self.configer.get('details', 'color_list')[(label - 1) % color_num], thickness=3) cv2.putText(ori_img, class_name, (int(rois[j][0]) + 5, int(rois[j][3]) - 5), cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=self.configer.get('details', 'color_list')[(label - 1) % color_num], thickness=2) img_path = os.path.join(base_dir, '{}_{}_{}.jpg'.format(name, i, time.time())) cv2.imwrite(img_path, ori_img)
def vis_peaks(self, heatmap_in, ori_img_in, name='default', sub_dir='peaks'): base_dir = os.path.join(self.configer.get('project_dir'), POSE_DIR, sub_dir) if not os.path.exists(base_dir): Log.error('Dir:{} not exists!'.format(base_dir)) os.makedirs(base_dir) if not isinstance(heatmap_in, np.ndarray): if len(heatmap_in.size()) != 3: Log.error('Heatmap size is not valid.') exit(1) heatmap = heatmap_in.clone().data.cpu().numpy().transpose(1, 2, 0) else: heatmap = heatmap_in.copy() if not isinstance(ori_img_in, np.ndarray): ori_img = DeNormalize(div_value=self.configer.get('normalize', 'div_value'), mean=self.configer.get('normalize', 'mean'), std=self.configer.get('normalize', 'std'))(ori_img_in.clone()) ori_img = ori_img.data.cpu().squeeze().numpy().transpose(1, 2, 0).astype(np.uint8) ori_img = cv2.cvtColor(ori_img, cv2.COLOR_RGB2BGR) else: ori_img = ori_img_in.copy() for j in range(self.configer.get('data', 'num_kpts')): peaks = self.__get_peaks(heatmap[:, :, j]) for peak in peaks: ori_img = cv2.circle(ori_img, (peak[0], peak[1]), self.configer.get('vis', 'circle_radius'), self.configer.get('details', 'color_list')[j], thickness=-1) cv2.imwrite(os.path.join(base_dir, '{}_{}.jpg'.format(name, j)), ori_img)
def vis_paf(self, inputs_in, ori_img_in, name='default', sub_dir='pafs'): base_dir = os.path.join(self.configer.get('project_dir'), POSE_DIR, sub_dir) if not os.path.exists(base_dir): Log.error('Dir:{} not exists!'.format(base_dir)) os.makedirs(base_dir) if not isinstance(inputs_in, np.ndarray): if len(inputs_in.size()) != 3: Log.error('Pafs size is not valid.') exit(1) inputs = inputs_in.clone().data.squeeze().cpu().numpy().transpose( 1, 2, 0) else: inputs = inputs_in.copy() if not isinstance(ori_img_in, np.ndarray): if len(ori_img_in.size()) != 3: Log.error('Image size is not valid.') exit(1) ori_img = DeNormalize( div_value=self.configer.get('normalize', 'div_value'), mean=self.configer.get('normalize', 'mean'), std=self.configer.get('normalize', 'std'))(ori_img_in.clone()) ori_img = ori_img.data.cpu().squeeze().numpy().transpose( 1, 2, 0).astype(np.uint8) else: ori_img = ori_img_in.copy() for i in range(len(self.configer.get('details', 'limb_seq'))): U = inputs[:, :, 2 * i] * -1 V = inputs[:, :, 2 * i + 1] X, Y = np.meshgrid(np.arange(U.shape[1]), np.arange(U.shape[0])) M = np.zeros(U.shape, dtype='bool') M[U**2 + V**2 < 0.5 * 0.5] = True U = ma.masked_array(U, mask=M) V = ma.masked_array(V, mask=M) # 1 img = Image.fromarray(ori_img, mode='RGB') plt.figure() plt.imshow(img, alpha=1.0) s = 5 Q = plt.quiver(X[::s, ::s], Y[::s, ::s], U[::s, ::s], V[::s, ::s], scale=50, headaxislength=4, alpha=1.0, width=0.001, color='r') fig = matplotlib.pyplot.gcf() fig.set_size_inches(10, 10) plt.savefig(os.path.join(base_dir, '{}_{}.jpg'.format(name, i)))
def vis_fn(self, preds, targets, ori_img_in=None, name='default', sub_dir='fn'): base_dir = os.path.join(self.configer.get('project_dir'), SEG_DIR, sub_dir) if not os.path.exists(base_dir): Log.error('Dir:{} not exists!'.format(base_dir)) os.makedirs(base_dir) if not isinstance(preds, np.ndarray): if len(preds.size()) > 3: Log.error('Preds size is not valid.') exit(1) if len(preds.size()) == 3: preds = preds.clone().data.cpu().numpy() if len(preds.size()) == 2: preds = preds.unsqueeze(0).data.cpu().numpy() else: if len(preds.shape) > 3: Log.error('Preds size is not valid.') exit(1) if len(preds.shape) == 2: preds = preds.unsqueeze(0) if not isinstance(targets, np.ndarray): if len(targets.size()) == 3: targets = targets.clone().data.cpu().numpy() if len(targets.size()) == 2: targets = targets.unsqueeze(0).data.cpu().numpy() else: if len(targets.shape) == 2: targets = targets.unsqueeze(0) if ori_img_in is not None: if not isinstance(ori_img_in, np.ndarray): if len(ori_img_in.size()) < 3: Log.error('Image size is not valid.') exit(1) if len(ori_img_in.size()) == 4: ori_img_in = ori_img_in.data.cpu() if len(ori_img_in.size()) == 3: ori_img_in = ori_img_in.unsqueeze(0).data.cpu() ori_img = ori_img_in.clone() for i in range(ori_img_in.size(0)): ori_img[i] = DeNormalize( div_value=self.configer.get('normalize', 'div_value'), mean=self.configer.get('normalize', 'mean'), std=self.configer.get('normalize', 'std'))(ori_img_in.clone()) ori_img = ori_img.numpy().transpose(2, 3, 1).astype(np.uint8) else: if len(ori_img_in.shape) == 3: ori_img_in = ori_img_in.unsqueeze(0) ori_img = ori_img_in.copy() for img_id in range(preds.shape[0]): label = targets[img_id] pred = preds[img_id] result = np.zeros(shape=(pred.shape[0], pred.shape[1], 3), dtype=np.uint8) for i in range(self.configer.get('data', 'num_classes')): mask0 = np.zeros_like(label, dtype=np.uint8) mask1 = np.zeros_like(label, dtype=np.uint8) mask0[label[:] == i] += 1 mask0[pred[:] == i] += 1 mask1[pred[:] == i] += 1 result[mask0[:] == 1] = self.configer.get( 'details', 'color_list')[i] result[mask1[:] == 1] = (0, 0, 0) image_result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB) if ori_img_in is not None: image_result = cv2.addWeighted(ori_img[i], 0.6, image_result, 0.4, 0) cv2.imwrite( os.path.join(base_dir, '{}_{}.jpg'.format(name, img_id)), image_result)