|
|
| import gradio as gr |
| from transformers import AutoImageProcessor, AutoModelForObjectDetection |
| import torch |
| from PIL import Image, ImageDraw |
|
|
| |
| |
| |
|
|
| |
| model_id = "daniihc16/yolo_finetuned_raccoon" |
|
|
| try: |
| image_processor = AutoImageProcessor.from_pretrained(model_id) |
| model = AutoModelForObjectDetection.from_pretrained(model_id) |
| except Exception as e: |
| print(f"Error cargando modelo: {e}. Asegúrate de poner el ID correcto.") |
| raise e |
|
|
| def predict(image): |
| if image is None: return None |
| |
| inputs = image_processor(images=image, return_tensors="pt") |
| |
| with torch.no_grad(): |
| outputs = model(**inputs) |
| |
| target_sizes = torch.tensor([image.size[::-1]]) |
| |
| results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[0] |
| |
| draw = ImageDraw.Draw(image) |
| |
| for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): |
| box = [round(i, 2) for i in box.tolist()] |
| x, y, x2, y2 = tuple(box) |
| |
| |
| draw.rectangle((x, y, x2, y2), outline="red", width=3) |
| |
| |
| label_name = model.config.id2label[label.item()] |
| draw.text((x, y), f"{label_name}: {round(score.item(), 2)}", fill="red") |
| |
| return image |
|
|
| iface = gr.Interface( |
| fn=predict, |
| inputs=gr.Image(type="pil"), |
| outputs=gr.Image(type="pil"), |
| title="Detector de Mapaches (Raccoon Detection)", |
| description="Sube una imagen para detectar mapaches usando un modelo YOLOS Finetuned.", |
| examples=['raccoon-1.jpg', 'raccoon-12.jpg'] |
| ) |
|
|
| if __name__ == "__main__": |
| iface.launch() |
|
|
