File size: 2,645 Bytes
2d0567e 6bdf257 2d0567e |
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 |
import numpy as np
import mediapipe as mp
from mediapipe import solutions
from mediapipe.framework.formats import landmark_pb2
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
# Crear un objeto PoseLandmarker
model_asset_path = 'models/pose_landmarker_heavy.task'
base_options = python.BaseOptions(model_asset_path, delegate=python.BaseOptions.Delegate.CPU)
def draw_landmarks_on_image(rgb_image, detection_result):
"""
Dibuja los puntos de referencia de la pose en la imagen.
Args:
rgb_image (np.ndarray): Imagen RGB de entrada.
detection_result: Resultado de la detección de pose.
Returns:
np.ndarray: Imagen anotada con los puntos de referencia de la pose.
"""
pose_landmarks_list = detection_result.pose_landmarks
annotated_image = np.copy(rgb_image)
# Recorrer las poses detectadas para visualizarlas
for pose_landmarks in pose_landmarks_list:
# Dibujar los puntos de referencia de la pose
pose_landmarks_proto = landmark_pb2.NormalizedLandmarkList()
pose_landmarks_proto.landmark.extend([
landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z) for landmark in pose_landmarks])
solutions.drawing_utils.draw_landmarks(
annotated_image,
pose_landmarks_proto,
solutions.pose.POSE_CONNECTIONS,
solutions.drawing_styles.get_default_pose_landmarks_style())
return annotated_image
def load_model(input_img, pos, confidence):
"""
Carga el modelo de detección de pose y lo aplica a la imagen de entrada.
Args:
input_img (np.ndarray): La imagen de entrada.
pos (float): Confianza mínima para la detección de poses.
confidence (int): Número máximo de poses a detectar.
Returns:
np.ndarray: Imagen anotada con los resultados de la detección de poses.
"""
# Configuración del objeto PoseLandmarker con parámetros própios
options = vision.PoseLandmarkerOptions(
base_options=base_options,
num_poses=confidence,
min_pose_detection_confidence=pos,
min_pose_presence_confidence=pos,
min_tracking_confidence=pos)
detector = vision.PoseLandmarker.create_from_options(options)
rgb_frame = mp.Image(image_format=mp.ImageFormat.SRGB, data=input_img)
# Detectar los puntos de referencia de la pose en la imagen de entrada
detection_result = detector.detect(rgb_frame)
# Procesar el resultado de la detección y visualizarlo
annotated_image = draw_landmarks_on_image(rgb_frame.numpy_view(), detection_result)
return annotated_image
|