Spaces:
Running
Running
File size: 3,051 Bytes
d015578 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
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
|