import cv2
import os
import torch
from facenet_pytorch import MTCNN
import argparse

# Função para processar o vídeo
def process_video(source, output_folder):
    # Verificar se a pasta de saída existe, se não, criar
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Carregar vídeo
    cap = cv2.VideoCapture(source)
    if not cap.isOpened():
        print(f"Erro ao abrir o vídeo {source}")
        return

    # Carregar o modelo MTCNN para detecção de rostos
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    mtcnn = MTCNN(keep_all=True, device=device)

    frame_count = 0
    face_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        frame_count += 1
        # Detecção de rostos
        boxes, _ = mtcnn.detect(frame)

        if boxes is not None:
            for box in boxes:
                # Cortar o rosto detectado
                x1, y1, x2, y2 = [int(coord) for coord in box]
                face = frame[y1:y2, x1:x2]

                # Salvar o rosto na pasta de saída
                face_filename = os.path.join(output_folder, f"face_{frame_count}_{face_count}.jpg")
                cv2.imwrite(face_filename, face)
                face_count += 1

        # Mostrar o progresso (opcional)
        print(f"Frame {frame_count} processado.")

    cap.release()
    print(f"Processamento concluído. {face_count} rostos salvos em {output_folder}.")

# Função principal para parsear argumentos e chamar o processamento
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Detecta e salva rostos de um vídeo.")
    parser.add_argument("--source", type=str, required=True, help="Caminho para o arquivo de vídeo.")
    parser.add_argument("--output", type=str, required=True, help="Pasta de saída para salvar os rostos detectados.")
    args = parser.parse_args()

    process_video(args.source, args.output)