app.py

import gradio as gr
import cv2
from ultralytics import YOLO
from supervision import Detections, ByteTrack

# Carregar o modelo YOLOv8 pré-treinado
model = YOLO("yolov8n.pt")

# Inicializar o rastreador ByteTrack
tracker = ByteTrack()

def process_video(video_file):
    """
    Função que processa um vídeo, realizando a detecção de objetos com YOLOv8
    e rastreamento com ByteTrack. O vídeo é então retornado com as detecções e
    IDs dos objetos desenhados nas caixas delimitadoras.
    """
    video_path = video_file.name
    cap = cv2.VideoCapture(video_path)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter("output.mp4", fourcc, 30.0, (int(cap.get(3)), int(cap.get(4))))

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

        # Detecção com YOLOv8
        results = model(frame)[0]

        # Converter as detecções para o formato necessário para o rastreador
        detections = Detections(
            xyxy=results.boxes.xyxy.cpu().numpy(),
            confidence=results.boxes.conf.cpu().numpy(),
            class_id=results.boxes.cls.cpu().numpy().astype(int)
        )

        # Aplicar rastreamento com ByteTrack
        tracks = tracker.update_with_detections(detections)

        # Desenhar as caixas de detecção e os IDs de rastreamento
        for track in tracks:
            box = track[0]
            x1, y1, x2, y2 = box
            track_id = track[4]
            class_id = track[3]
            label = f"{model.names[class_id]} {track_id}"

            # Desenhar a caixa delimitadora e o ID do rastreamento
            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
            cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        out.write(frame)

    cap.release()
    out.release()

    return "output.mp4"

# Criando uma interface de usuário com Gradio
iface = gr.Interface(
    fn=process_video,
    inputs=gr.Video(label="Envie um vídeo para rastreamento"),
    outputs="video",
    title="Aplicação de Rastreamento de Objetos com YOLOv8 e ByteTrack",
    description="""
    Este aplicativo utiliza um modelo YOLOv8 para detectar objetos em vídeos e
    um rastreador ByteTrack para acompanhar o movimento dos objetos ao longo do tempo.
    A cada quadro do vídeo, as detecções de objetos são desenhadas com caixas delimitadoras
    e IDs de rastreamento.

    **Como funciona:**
    1. Envie um vídeo.
    2. O modelo YOLOv8 detecta objetos em cada quadro.
    3. O rastreador ByteTrack segue os objetos detectados.
    4. O vídeo de saída é gerado com as caixas delimitadoras e os IDs de rastreamento desenhados.

    **Exemplo:**
    Veja o exemplo de como o sistema funciona no GIF abaixo!
    """,
    examples=["example.gif"],  # Adicionando um GIF de exemplo
)

# Executar a interface
iface.launch()