174 lines
7.6 KiB
Python
174 lines
7.6 KiB
Python
# MIT License
|
|
#
|
|
# Copyright (c) 2019-2024 Iván de Paz Centeno
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in all
|
|
# copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
# SOFTWARE.
|
|
|
|
import numpy as np
|
|
|
|
from .bboxes import parse_bbox
|
|
from .landmarks import parse_landmarks
|
|
|
|
|
|
def plot_bbox(image, bbox, color="#FFFF00", normalize_color=False, input_as_width_height=True):
|
|
"""
|
|
Draws a bounding box on the given image.
|
|
|
|
Args:
|
|
image (np.ndarray): The input image on which to draw the bounding box.
|
|
bbox (list, dict, or np.ndarray): The bounding box coordinates in one of the following formats:
|
|
- List or array [x1, y1, x2, y2].
|
|
- Dict with a key 'box' that contains the bounding box.
|
|
color (str): The color of the bounding box in hex format (default is yellow, "#FFFF00").
|
|
normalize_color (bool): True if color should be in [0..1]. False to make it between [0..255]
|
|
input_as_width_height (bool): True if input `bbox` parameter follows format [x1, y1, width, height]. False if follows format [x1, y1, x2, y2]
|
|
Returns:
|
|
np.ndarray: The image with the bounding box drawn.
|
|
"""
|
|
color = parse_color(color) # Convert color to RGB
|
|
color = color if normalize_color else (color * 255).astype(np.uint8)
|
|
|
|
# Parse the bounding box coordinates
|
|
bbox = parse_bbox(bbox, input_as_width_height=input_as_width_height, output_as_width_height=False)
|
|
|
|
image = image.copy() # Copy the image to avoid modifying the original
|
|
|
|
# Draw the vertical sides of the box (left and right)
|
|
image[bbox[1]:bbox[3], bbox[0], :] = color # Left side
|
|
image[bbox[1]:bbox[3], bbox[2], :] = color # Right side
|
|
|
|
# Draw the horizontal sides of the box (top and bottom)
|
|
image[bbox[1], bbox[0]:bbox[2], :] = color # Top side
|
|
image[bbox[3], bbox[0]:bbox[2], :] = color # Bottom side
|
|
|
|
result = image if normalize_color else image.astype(np.uint8)
|
|
return result
|
|
|
|
|
|
def plot_landmarks(image, landmarks, color="#FFFF00", keypoints="nose,mouth_right,right_eye,left_eye,mouth_left",
|
|
brush_size=2, normalize_color=False):
|
|
"""
|
|
Plots facial landmarks on the given image.
|
|
|
|
Args:
|
|
image (np.ndarray): The input image on which to draw the landmarks.
|
|
landmarks (dict or np.ndarray): The facial landmarks to plot, either as a dictionary or an array.
|
|
color (str): The color of the landmarks in hex format (default is yellow, "#FFFF00").
|
|
keypoints (str): A comma-separated list of keypoints to plot (default includes all facial landmarks).
|
|
brush_size (int): The size of the brush used to draw the keypoints (default is 2).
|
|
normalize_color (bool): True if color should be in [0..1]. False to make it between [0..255]
|
|
|
|
Returns:
|
|
np.ndarray: The image with the landmarks drawn.
|
|
"""
|
|
keypoints = [k.strip() for k in keypoints.split(",")] # Parse the keypoints list
|
|
color = parse_color(color) # Convert color to RGB
|
|
color = color if normalize_color else (color * 255).astype(np.uint8)
|
|
|
|
try:
|
|
landmarks = parse_landmarks(landmarks) # Parse the landmarks
|
|
|
|
except IndexError: # No landmarks available
|
|
return image
|
|
|
|
image = image.copy() # Copy the image to avoid modifying the original
|
|
|
|
# Draw each landmark as a small circle
|
|
for key in keypoints:
|
|
if key in landmarks:
|
|
x, y = landmarks[key]
|
|
image[y-brush_size:y+brush_size+1, x-brush_size:x+brush_size+1, :] = color # Draw the landmark
|
|
|
|
result = image if normalize_color else image.astype(np.uint8)
|
|
return result
|
|
|
|
|
|
def plot(image, detection, input_as_width_height=True):
|
|
"""
|
|
Plots a single or multiple facial detection results on the given image.
|
|
|
|
Args:
|
|
image (np.ndarray): The input image on which to draw the detections.
|
|
detection (list, dict, or np.ndarray): A single detection or a list/array of detections to plot.
|
|
Each detection contains facial landmarks and/or bounding box information.
|
|
input_as_width_height (bool): Whether the input bounding box format is (width, height) instead of
|
|
the default (x1, y1, x2, y2) (default is True).
|
|
|
|
Returns:
|
|
np.ndarray or None: The image with the detection(s) plotted, or None if no detection is present.
|
|
"""
|
|
if len(detection) == 0:
|
|
return None
|
|
|
|
if isinstance(detection, list) or (isinstance(detection, np.ndarray) and len(detection.shape) > 1):
|
|
return plot_all(image, detection, input_as_width_height=input_as_width_height)
|
|
|
|
return plot_landmarks(plot_bbox(image, detection, input_as_width_height=input_as_width_height), detection)
|
|
|
|
|
|
def plot_all(image, detections, input_as_width_height=True):
|
|
"""
|
|
Plots multiple facial detection results on the given image.
|
|
|
|
Args:
|
|
image (np.ndarray): The input image on which to draw the detections.
|
|
detections (list or np.ndarray): A list or array of detections, where each detection contains
|
|
facial landmarks and/or bounding box information.
|
|
input_as_width_height (bool): Whether the input bounding box format is (width, height) instead of
|
|
the default (x1, y1, x2, y2) (default is True).
|
|
|
|
Returns:
|
|
np.ndarray: The image with all detections plotted.
|
|
"""
|
|
for detection in detections:
|
|
image = plot_landmarks(plot_bbox(image, detection, input_as_width_height=input_as_width_height), detection)
|
|
|
|
return image
|
|
|
|
|
|
def parse_color(color):
|
|
"""
|
|
Parses a color from a string in various formats (e.g., hex, RGB) into a normalized RGB array.
|
|
|
|
The color can be provided in the following formats:
|
|
* Hexadecimal string (e.g., "#RRGGBB" or "#RGB")
|
|
* Hexadecimal string with prefix "0x" (e.g., "0xRRGGBB")
|
|
|
|
Args:
|
|
color (str): A color string in hex format (e.g., "#RRGGBB", "#RGB", "0xRRGGBB").
|
|
|
|
Returns:
|
|
np.ndarray: A numpy array of normalized RGB values (between 0 and 1) representing the color.
|
|
"""
|
|
if isinstance(color, str):
|
|
if color.startswith("#"):
|
|
color = color[1:] # Remove '#' prefix
|
|
if color.startswith("0x"):
|
|
color = color[2:] # Remove '0x' prefix
|
|
if len(color) == 3: # Short form hex color (#RGB)
|
|
color = np.asarray([int(f"{color[0]}{color[0]}", base=16),
|
|
int(f"{color[1]}{color[1]}", base=16),
|
|
int(f"{color[2]}{color[2]}", base=16)]) / 255
|
|
if len(color) == 6: # Full form hex color (#RRGGBB)
|
|
color = np.asarray([int(f"{color[0]}{color[1]}", base=16),
|
|
int(f"{color[2]}{color[3]}", base=16),
|
|
int(f"{color[4]}{color[5]}", base=16)]) / 255
|
|
|
|
return color
|