VARIOS – Visión: múltiple pose de personas con YOLOv7

Para detectar más de una persona con su pose en una imagen no vale MediaPipe y hay que usar OpenPose, PoseNET o YOLO. En esta prueba comento cómo conseguir la detección múltiple de personas en una vista de cámara y sin necesidad de tarjeta GPU, aunque a muy baja velocidad.

Para esta prueba me he basado en el ejemplo de LearnOpenCV, en el siguiente enlace:

https://github.com/spmallick/learnopencv/tree/master/YOLOv7-Pose-vs-MediaPipe-in-Human-Pose-Estimation

En la carpeta Experiments se encuentra el código en el fichero yolo-v7-pose.py. Para conseguir que funcione hay que instalar una serie de librerías y otros requisitos. El ejemplo que comento lo he hecho con Python 3.8.10, aunque supongo que con otras versiones de Python también puede funcionar.

Requisitos

• Instalar las siguientes librerías con pip: opencv-python, tqdm, pyyaml, seaborn.
• Instalar PyTorch: esta librería es importante y delicada y no se puede instalar a través de pip. En mi caso la he instalado para Windows sin tarjeta gráfica Nvidia, lo que quiere decir sin CUDA. En el siguiente enlace se detalla cómo instalarlo en cada caso: https://www.geeksforgeeks.org/install-pytorch-on-windows/
• Hay que copiar algunos ficheros de YOLOv7. En el siguiente enlace se puede descargar el sistema completo oficial, con los ficheros de entrenamiento: https://github.com/WongKinYiu/yolov7. No es necesario copiarlo todo, solamente son necesarias las carpetas models/ y utils/.
• Copiar el fichero de entrenamiento de poses, con los pesos de la red neuronal. Este fichero se encuentra en la página que pongo en el apartado anterior pero no en la estructura general de ficheros sino en un enlace al final de la misma, que reproduzco aquí: https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-w6-pose.pt

Haciendo lo anterior ya debería funcionar. El código modificado es el siguiente:

import time
import torch
import cv2
import numpy as np
from torchvision import transforms

from utils.datasets import letterbox
from utils.general  import non_max_suppression_kpt
from utils.plots    import output_to_keypoint, plot_skeleton_kpts


def pose_video(frame):
    mapped_img = frame.copy()
    # Letterbox resizing.
    img = letterbox(frame, input_size, stride=64, auto=True)[0]
    #print(img.shape)
    img_ = img.copy()
    # Convert the array to 4D.
    img = transforms.ToTensor()(img)
    # Convert the array to Tensor.
    img = torch.tensor(np.array([img.numpy()]))
    # Load the image into the computation device.
    img = img.to(device)
    
    # Gradients are stored during training, not required while inference.
    with torch.no_grad():
        t1 = time.time()
        output, _ = model(img)
        t2 = time.time()
        fps = 1/(t2 - t1)
        output = non_max_suppression_kpt(output, 
                                         0.25,    # Conf. Threshold.
                                         0.65,    # IoU Threshold.
                                         nc=1,   # Number of classes.
                                         nkpt=17, # Number of keypoints.
                                         kpt_label=True)
        
        output = output_to_keypoint(output)

    # Change format [b, c, h, w] to [h, w, c] for displaying the image.
    nimg = img[0].permute(1, 2, 0) * 255
    nimg = nimg.cpu().numpy().astype(np.uint8)
    nimg = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)

    for idx in range(output.shape[0]):
        plot_skeleton_kpts(nimg, mapped_img, input_size, output[idx, 7:].T, 3)
        
    return nimg, fps


# Change forward pass input size.
input_size = 512
device = torch.device("cpu")

# Load keypoint detection model.
weights = torch.load('yolov7-w6-pose.pt', map_location=device)
model = weights['model']
# Load the model in evaluation mode.
_ = model.float().eval()
# Load the model to computation device [cpu/gpu/tpu]
model.to(device)

cap = cv2.VideoCapture(0)
fps = int(cap.get(cv2.CAP_PROP_FPS))
ret, frame = cap.read()
h, w, _ = frame.shape

if __name__ == '__main__':
        while True:
            ret, frame = cap.read()
            
            if not ret:
                print('Unable to read frame. Exiting ..')
                break

            img, fps_ = pose_video(frame)

            cv2.putText(img, 'FPS : {:.2f}'.format(fps_), (120, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2, cv2.LINE_AA)
            cv2.putText(img, 'YOLOv7', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2, cv2.LINE_AA)

            cv2.imshow('Output', img[...,::-1])
            if (cv2.waitKey(1) == ord('s')):
                break

        cap.release()
        cv2.destroyAllWindows()

En este código de ejemplo he puesto el tamaño de la imagen a 512 pixeles, pero se podría cambiar. El resultado consigue muy pocas imágenes por segundo, menos de 1. Para mejorar este dato habría que usa una tarjeta aceleradora. A continuación una imagen de prueba.