Spaces:
Sleeping
Sleeping
File size: 2,762 Bytes
b2ab75b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import numpy as np
import supervision as sv
from ultralytics import YOLO
import gradio as gr
import torch
import spaces
import os
tracker = sv.ByteTrack()
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
trace_annotator = sv.TraceAnnotator()
device = "cuda" if torch.cuda.is_available() else "cpu"
model = YOLO("yolov8n.pt").to(device)
FRAME_INTERVAL = 35
LOW_RES_SIZE = (640,640)
MID_RES_SIZE = (1280,1280)
HIGH_RES_SIZE =(1088,1920)
MAX_PERSONS = 0
auth_users = [(os.getenv('USERNAME'), os.getenv('PASSWORD'))]
@spaces.GPU
def process_video(frame: np.ndarray, _: int) -> np.ndarray:
global MAX_PERSONS
results = model(frame, imgsz=HIGH_RES_SIZE)[0]
detections = sv.Detections.from_ultralytics(results)
person_detections = detections[detections.class_id == 0]
current_person_count = len(person_detections)
if current_person_count > MAX_PERSONS:
MAX_PERSONS = current_person_count
print(f'Personas detectadas: {current_person_count}, Máximo: {MAX_PERSONS}')
tracked_persons = tracker.update_with_detections(person_detections)
labels = [
f"#{tracker_id} {results.names[class_id]}"
for class_id, tracker_id
in zip(tracked_persons.class_id, tracked_persons.tracker_id)
]
annotated_frame = box_annotator.annotate(frame.copy(), detections=tracked_persons)
annotated_frame = label_annotator.annotate(annotated_frame, detections=tracked_persons, labels=labels)
return trace_annotator.annotate(annotated_frame, detections=tracked_persons)
def upload_video(video_path):
global MAX_PERSONS
MAX_PERSONS = 0 # Reiniciar el contador para cada video
output_video_path = "output_video.mp4"
if os.path.exists(output_video_path):
os.remove(output_video_path)
sv.process_video(source_path=video_path, target_path=output_video_path, callback=process_video)
return output_video_path, MAX_PERSONS
with gr.Blocks() as demo:
gr.Markdown("# Aegis Air Demo")
with gr.Row():
video_input = gr.Video(label="Sube tu video aquí")
video_output = gr.Video(label="Video Anotado")
persons_output = gr.Textbox(label="Total de Personas Detectadas", interactive=False)
def process_video_gradio(video):
if video is None:
return None, "No se ha subido ningún video"
processed_video, total_persons = upload_video(video)
return processed_video, f"Máximo número de personas detectadas: {total_persons}"
submit_button = gr.Button("Procesar")
submit_button.click(process_video_gradio, inputs=video_input, outputs=[video_output, persons_output])
demo.launch(auth=auth_users) |