Video_processor / app.py
User00100's picture
Create app.py
e90b086 verified
raw
history blame
1.02 kB
import gradio as gr
import cv2
import torch
from transformers import DetrImageProcessor, DetrForObjectDetection
# Cargar el modelo y el procesador
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
def detect_objects(video_path):
# Leer el video
cap = cv2.VideoCapture(video_path)
results = []
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# Procesar el frame
inputs = processor(images=frame, return_tensors="pt")
outputs = model(**inputs)
# Extraer las predicciones
target_sizes = torch.tensor([frame.shape[:2]])
results.append(processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0])
cap.release()
return results
# Crear la interfaz de Gradio
iface = gr.Interface(fn=detect_objects, inputs="video", outputs="json")
iface.launch()