Spaces:
Running
Running
import matplotlib.pyplot as plt | |
import numpy as np | |
import cv2 | |
import spiga.data.loaders.augmentors.utils as dlu | |
BLUE = (255, 0, 0) | |
GREEN = (0, 255, 0) | |
RED = (0, 0, 255) | |
PURPLE = (128, 0, 128) | |
def draw_landmarks(image, landmarks, visible=None, mask=None, thick_scale=1, colors=(GREEN, RED)): | |
# Fix variable | |
thick = int(2 * thick_scale + 0.5) | |
# Initialize variables if need it | |
if visible is None: | |
visible = np.ones(len(landmarks)) | |
if mask is None: | |
mask = np.ones(len(landmarks)) | |
mask = np.array(mask, dtype=bool) | |
visible = np.array(visible, dtype=bool) | |
# Clean and split landmarks | |
landmarks = landmarks[mask] | |
visible = visible[mask] | |
ldm_vis = landmarks[visible] | |
not_visible = np.logical_not(visible) | |
ldm_notvis = landmarks[not_visible] | |
# Plot landmarks | |
if image.shape[0] == 3: | |
image = image.transpose(1, 2, 0) | |
canvas = image.copy() | |
canvas = _write_circles(canvas, ldm_vis, color=colors[0], thick=thick) | |
canvas = _write_circles(canvas, ldm_notvis, color=colors[1], thick=thick) | |
return canvas | |
def _write_circles(canvas, landmarks, color=RED, thick=2): | |
for xy in landmarks: | |
xy = np.array(xy+0.5, dtype=int) | |
canvas = cv2.circle(canvas, (xy[0], xy[1]), thick, color, -1) | |
return canvas | |
def plot_landmarks_pil(image, landmarks, visible=None, mask=None): | |
# Initialize variables if need it | |
if visible is None: | |
visible = np.ones(len(landmarks)) | |
if mask is None: | |
mask = np.ones(len(landmarks)) | |
mask = np.array(mask, dtype=bool) | |
visible = np.array(visible, dtype=bool) | |
not_visible = np.logical_not(visible) | |
# Clean and split landmarks | |
landmarks = landmarks[mask] | |
ldm_vis = landmarks[visible] | |
ldm_notvis = landmarks[not_visible] | |
# Plot landmarks | |
if image.shape[0] == 3: | |
image = image.transpose(1, 2, 0) | |
plt.imshow(image / 255) | |
plt.scatter(ldm_vis[:, 0], ldm_vis[:, 1], s=10, marker='.', c='g') | |
plt.scatter(ldm_notvis[:, 0], ldm_notvis[:, 1], s=10, marker='.', c='r') | |
plt.show() | |
def draw_pose(img, rot, trl, K, euler=False, size=0.5, colors=(BLUE, GREEN, RED)): | |
if euler: | |
rot = dlu.euler_to_rotation_matrix(rot) | |
canvas = img.copy() | |
rotV, _ = cv2.Rodrigues(rot) | |
points = np.float32([[size, 0, 0], [0, -size, 0], [0, 0, -size], [0, 0, 0]]).reshape(-1, 3) | |
axisPoints, _ = cv2.projectPoints(points, rotV, trl, K, (0, 0, 0, 0)) | |
axisPoints = axisPoints.astype(int) | |
canvas = cv2.line(canvas, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), colors[0], 3) | |
canvas = cv2.line(canvas, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), colors[1], 3) | |
canvas = cv2.line(canvas, tuple(axisPoints[3].ravel()), tuple(axisPoints[0].ravel()), colors[2], 3) | |
return canvas | |
def enhance_heatmap(heatmap): | |
map_aux = heatmap - heatmap.min() | |
map_aux = map_aux / map_aux.max() | |
map_img = cv2.applyColorMap((map_aux * 255).astype(np.uint8), cv2.COLORMAP_BONE) | |
return map_img | |