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