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

En esta entrada voy a contar mi experiencia con la detección de poses usando un ordenador con tarjeta aceleradora. El entorno es el siguiente:

• Ordenador con procesador i7 y 16 GB de memoria RAM.
• Tarjeta gráfica NVidia RTX3070 con 8 GB de memoria.
• Sistema operativo Windows 11.
• CUDA Toolkit 12.0.

Voy a seguir más o menos los mismos pasos que en el artículo sobre poses con YOLOv7 con CPU. Para esta prueba me he basado originalmente en el ejemplo de LearnOpenCV, en el siguiente enlace:

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

Requisitos

• Instalar las siguientes librerías con pip: opencv-python, tqdm, pyyaml, seaborn, scipy.
• 
• 
• Instalar PyTorch: esta librería es importante y delicada y no se puede instalar simplemente con un pip. Para instalarla con éxito se hace desde la página oficial: https://pytorch.org/get-started/locally/ En esta página he seleccionado las opciones: Stable, Windows, pip, Python, CUDA 11.7. Aparece en el cuadro inferior el comando que hay que ejecutar, que se copia y pega en una ventana de CMD, y tras descargarse la librería y otras dependencias, me ha funcionado bien.
• 
• 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 para este caso 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, output[idx, 7:].T, 3)
        
    return nimg, fps


# Change forward pass input size.
input_size = 512
device = torch.device("cuda:0")

# 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 también el tamaño de la imagen a 512 pixeles, pero se podría cambiar. El resultado consigue alrededor de 15 imágenes por segundo, unas 25 veces el resultado sin GPU, lo que no está mal, aunque no es todavía vídeo fluido a tiempo real. A continuación una imagen de prueba.