daniihc16 commited on
Commit
57d100e
·
verified ·
1 Parent(s): 50278f7

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +59 -0
  2. requirements.txt +5 -0
app.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import gradio as gr
3
+ from transformers import AutoImageProcessor, AutoModelForObjectDetection
4
+ import torch
5
+ from PIL import Image, ImageDraw
6
+
7
+ # Cargar modelo desde el Hub (Recomendado) o Local
8
+ # Si subiste tu modelo con trainer.push_to_hub(), usa tu ID: NO 'yolo_finetuned_raccoon' local.
9
+ # Ejemplo: model_id = "daniihc16/yolo_finetuned_raccoon" (Sustituye por tu usuario)
10
+
11
+ # Para facilitar la prueba, intentaremos cargar de una carpeta local si existe, sino del hub (si se configura).
12
+ # AQUÍ DEBES PONER EL ID DE TU MODELO SUBIDO A HUGGINGFACE
13
+ model_id = "hustvl/yolos-tiny" # Placeholder! CAMBIALO POR TU MODELO FINETUNED
14
+
15
+ try:
16
+ image_processor = AutoImageProcessor.from_pretrained(model_id)
17
+ model = AutoModelForObjectDetection.from_pretrained(model_id)
18
+ except Exception as e:
19
+ print(f"Error cargando modelo: {e}. Asegúrate de poner el ID correcto.")
20
+ raise e
21
+
22
+ def predict(image):
23
+ if image is None: return None
24
+
25
+ inputs = image_processor(images=image, return_tensors="pt")
26
+
27
+ with torch.no_grad():
28
+ outputs = model(**inputs)
29
+
30
+ target_sizes = torch.tensor([image.size[::-1]])
31
+ # Usamos un umbral de 0.5 para mostrar solo detecciones firmes
32
+ results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[0]
33
+
34
+ draw = ImageDraw.Draw(image)
35
+
36
+ for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
37
+ box = [round(i, 2) for i in box.tolist()]
38
+ x, y, x2, y2 = tuple(box)
39
+
40
+ # Dibujar caja
41
+ draw.rectangle((x, y, x2, y2), outline="red", width=3)
42
+
43
+ # Dibujar etiqueta
44
+ label_name = model.config.id2label[label.item()]
45
+ draw.text((x, y), f"{label_name}: {round(score.item(), 2)}", fill="red")
46
+
47
+ return image
48
+
49
+ iface = gr.Interface(
50
+ fn=predict,
51
+ inputs=gr.Image(type="pil"),
52
+ outputs=gr.Image(type="pil"),
53
+ title="Detector de Mapaches (Raccoon Detection)",
54
+ description="Sube una imagen para detectar mapaches usando un modelo YOLOS Finetuned.",
55
+ examples=[]
56
+ )
57
+
58
+ if __name__ == "__main__":
59
+ iface.launch()
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+
2
+ transformers
3
+ torch
4
+ pillow
5
+ gradio