Segizu commited on
Commit
27e96af
1 Parent(s): daa16ff
Files changed (1) hide show
  1. app.py +11 -30
app.py CHANGED
@@ -3,29 +3,9 @@ import gradio as gr
3
  from transformers import pipeline
4
  from PIL import Image
5
  import tempfile
6
- import torch
7
 
8
- device = "cuda" if torch.cuda.is_available() else "cpu"
9
- print(f"Using device: {device}")
10
-
11
- # Cargar el modelo de detecci贸n de objetos
12
- try:
13
- detector = pipeline(
14
- "object-detection",
15
- model="facebook/detr-resnet-50",
16
- device=0 if device == "cuda" else -1, # 0 para GPU, -1 para CPU
17
- framework="pt" # Especificar PyTorch como framework
18
- )
19
- print("Model loaded successfully on", device)
20
- except Exception as e:
21
- print(f"Error loading model: {e}")
22
- print("Falling back to CPU")
23
- detector = pipeline(
24
- "object-detection",
25
- model="facebook/detr-resnet-50",
26
- device=-1,
27
- framework="pt"
28
- )
29
 
30
  def process_video(video_path):
31
  """
@@ -47,11 +27,11 @@ def process_video(video_path):
47
  output_path = tmp_file.name
48
  tmp_file.close() # Se cierra para que VideoWriter pueda escribir en 茅l
49
 
50
- # Configurar VideoWriter (utilizamos el c贸dec mp4v)
51
  fourcc = cv2.VideoWriter_fourcc(*'mp4v')
52
  out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
53
 
54
- # Definir las clases a las que queremos aplicar detecci贸n
55
  valid_labels = {"person", "bicycle", "motorcycle"}
56
  threshold = 0.7 # Umbral de confianza
57
 
@@ -60,7 +40,7 @@ def process_video(video_path):
60
  if not ret:
61
  break
62
 
63
- # Convertir el frame de BGR a RGB y luego a imagen PIL
64
  frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
65
  pil_image = Image.fromarray(frame_rgb)
66
 
@@ -74,11 +54,12 @@ def process_video(video_path):
74
  if score < threshold or label not in valid_labels:
75
  continue
76
 
77
- # Obtener la caja del objeto en formato [xmin, ymin, width, height]
78
  box = detection["box"]
79
- xmin, ymin, w, h = box
80
- xmax = xmin + w
81
- ymax = ymin + h
 
82
 
83
  # Dibujar el rect谩ngulo y la etiqueta en el frame
84
  cv2.rectangle(frame, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color=(0, 255, 0), thickness=2)
@@ -97,7 +78,7 @@ iface = gr.Interface(
97
  inputs=gr.Video(label="Sube tu video"),
98
  outputs=gr.Video(label="Video procesado"),
99
  title="Detecci贸n y Visualizaci贸n de Objetos en Video",
100
- description="Carga un video y se detectan personas, bicicletas y motos. Los objetos se enmarcan y etiquetan, mostrando la detecci贸n en tiempo real."
101
  )
102
 
103
  if __name__ == "__main__":
 
3
  from transformers import pipeline
4
  from PIL import Image
5
  import tempfile
 
6
 
7
+ # Cargar el modelo de detecci贸n de objetos usando CPU
8
+ detector = pipeline("object-detection", model="facebook/detr-resnet-50", device=-1)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9
 
10
  def process_video(video_path):
11
  """
 
27
  output_path = tmp_file.name
28
  tmp_file.close() # Se cierra para que VideoWriter pueda escribir en 茅l
29
 
30
+ # Configurar VideoWriter (usamos el c贸dec mp4v)
31
  fourcc = cv2.VideoWriter_fourcc(*'mp4v')
32
  out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
33
 
34
+ # Definir las clases de inter茅s
35
  valid_labels = {"person", "bicycle", "motorcycle"}
36
  threshold = 0.7 # Umbral de confianza
37
 
 
40
  if not ret:
41
  break
42
 
43
+ # Convertir el frame de BGR a RGB y a imagen PIL
44
  frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
45
  pil_image = Image.fromarray(frame_rgb)
46
 
 
54
  if score < threshold or label not in valid_labels:
55
  continue
56
 
57
+ # Extraer la caja del objeto (dado que es un diccionario)
58
  box = detection["box"]
59
+ xmin = box["xmin"]
60
+ ymin = box["ymin"]
61
+ xmax = box["xmax"]
62
+ ymax = box["ymax"]
63
 
64
  # Dibujar el rect谩ngulo y la etiqueta en el frame
65
  cv2.rectangle(frame, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color=(0, 255, 0), thickness=2)
 
78
  inputs=gr.Video(label="Sube tu video"),
79
  outputs=gr.Video(label="Video procesado"),
80
  title="Detecci贸n y Visualizaci贸n de Objetos en Video",
81
+ description="Carga un video y se detectan personas, bicicletas y motos. Los objetos se enmarcan y etiquetan en tiempo real."
82
  )
83
 
84
  if __name__ == "__main__":