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_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_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_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)