Face detection using MTCNN

mtcnn_detection.py

import cv2
from facenet_pytorch.models.mtcnn import MTCNN
import numpy as np
import torch

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
mtcnn = MTCNN(keep_all=True, device=device)


def mtcnn_detect(img: np.ndarray) -> np.ndarray:
    boxes, probs = mtcnn.detect(img)
    for box in boxes:
        x_left = min(box[0], box[2])
        x_right = max(box[0], box[2])
        y_left = min(box[1], box[3])
        y_right = max(box[1], box[3])
        img = cv2.rectangle(img, (x_left, y_left), (x_right, y_right), 
                            (255, 0, 0), 2)

    return img