Update app.py

app.py

@@ -1,7 +1,20 @@
 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):
-    # Código de detecção e rastreamento adaptado para trabalhar com o vídeo enviado pelo usuário
+    """
+    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')
@@ -15,17 +28,17 @@ def process_video(video_file):
         # Detecção com YOLOv8
         results = model(frame)[0]
 
-        # Converter detecções para o formato do rastreador
+        # 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
+        # Aplicar rastreamento com ByteTrack
         tracks = tracker.update_with_detections(detections)
 
-        # Desenhar detecções e rastreamento
+        # Desenhar as caixas de detecção e os IDs de rastreamento
         for track in tracks:
             box = track[0]
             x1, y1, x2, y2 = box
@@ -33,6 +46,7 @@ def process_video(video_file):
             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)
 
@@ -40,7 +54,32 @@ def process_video(video_file):
 
     cap.release()
     out.release()
+
     return "output.mp4"
 
-iface = gr.Interface(fn=process_video, inputs=gr.Video(), outputs="video")
+# 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()