60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
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)
|