Spaces:
Running on Zero
Running on Zero
| import os | |
| import io | |
| import uuid | |
| import spaces | |
| import torch | |
| import torch.nn.functional as F | |
| import gradio as gr | |
| from PIL import Image | |
| from ultralytics import YOLO | |
| from transformers import AutoImageProcessor, AutoModel | |
| from qdrant_client import QdrantClient | |
| from qdrant_client.models import Distance, VectorParams, PointStruct | |
| import cloudinary | |
| import cloudinary.uploader | |
| from typing import List | |
| # ==================== 1. CONFIGURACIÓN DE CREDENCIALES ==================== | |
| QDRANT_URL = os.environ.get("QDRANT_URL") | |
| QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY") | |
| cloudinary.config( | |
| cloud_name=os.environ.get("CLOUDINARY_CLOUD_NAME"), | |
| api_key=os.environ.get("CLOUDINARY_API_KEY"), | |
| api_secret=os.environ.get("CLOUDINARY_API_SECRET"), | |
| secure=True | |
| ) | |
| # Conectar a Qdrant y crear la colección si no existe | |
| qdrant = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY) | |
| COLLECTION_NAME = "mascotas_perdidas" | |
| VECTOR_SIZE = 1024 # Dimensión exacta de facebook/convnext-base-224 | |
| if not qdrant.collection_exists(COLLECTION_NAME): | |
| qdrant.create_collection( | |
| collection_name=COLLECTION_NAME, | |
| vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE) | |
| ) | |
| # ==================== 2. CARGAR MODELOS DE IA ==================== | |
| print("Cargando YOLOv8 y ConvNeXt...") | |
| detector = YOLO("yolov8n.pt") | |
| model_name = "facebook/convnext-base-224" | |
| processor = AutoImageProcessor.from_pretrained(model_name) | |
| model = AutoModel.from_pretrained(model_name) | |
| model.eval() | |
| print("¡Modelos listos!") | |
| # ==================== 3. LÓGICA DE VISIÓN Y EMBEDDINGS (CON ZEROGPU) ==================== | |
| def crop_pet(pil_img: Image.Image) -> Image.Image: | |
| """Detecta y recorta el perro o gato de la imagen.""" | |
| results = detector(pil_img, verbose=False) | |
| PET_CLASSES = [15, 16] # 15: gato, 16: perro | |
| for r in results: | |
| for box in r.boxes: | |
| if int(box.cls[0]) in PET_CLASSES: | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| w, h = pil_img.size | |
| x1, y1 = max(0, x1 - 10), max(0, y1 - 10) | |
| x2, y2 = min(w, x2 + 10), min(h, y2 + 10) | |
| return pil_img.crop((x1, y1, x2, y2)) | |
| return pil_img | |
| def get_average_embedding(pil_images: List[Image.Image]) -> List[float]: | |
| """Genera el ADN visual promedio usando ZeroGPU.""" | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model.to(device) | |
| embeddings = [] | |
| for img in pil_images: | |
| cropped = crop_pet(img.convert("RGB")) | |
| inputs = processor(images=cropped, return_tensors="pt").to(device) | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| tensor = outputs.last_hidden_state.mean(dim=[2, 3]) | |
| embeddings.append(F.normalize(tensor, p=2, dim=1)) | |
| mean_emb = torch.mean(torch.stack(embeddings), dim=0) | |
| final_emb = F.normalize(mean_emb, p=2, dim=1) | |
| return final_emb.squeeze().cpu().tolist() | |
| # ==================== 4. FUNCIONES (INTERFAZ + API REST AUTOMÁTICA) ==================== | |
| def registrar_mascota(nombre: str, contacto: str, descripcion: str, archivos): | |
| if not archivos: | |
| return "⚠️ Sube entre 1 y 3 fotos del perro." | |
| imgs = [Image.open(f.name if hasattr(f, "name") else f) for f in archivos[:3]] | |
| # 1. Subir la foto principal a Cloudinary | |
| buf = io.BytesIO() | |
| imgs[0].convert("RGB").save(buf, format="JPEG", quality=85) | |
| buf.seek(0) | |
| upload_res = cloudinary.uploader.upload(buf, folder="mascotas_perdidas") | |
| foto_url = upload_res["secure_url"] | |
| # 2. Calcular embedding promedio (soporta cachorro + adulto) | |
| vector = get_average_embedding(imgs) | |
| # 3. Guardar en Qdrant Cloud | |
| pet_id = str(uuid.uuid4()) | |
| qdrant.upsert( | |
| collection_name=COLLECTION_NAME, | |
| points=[ | |
| PointStruct( | |
| id=pet_id, | |
| vector=vector, | |
| payload={ | |
| "nombre": nombre, | |
| "contacto": contacto, | |
| "descripcion": descripcion, | |
| "foto_url": foto_url | |
| } | |
| ) | |
| ] | |
| ) | |
| return f"✅ ¡Mascota guardada en Qdrant y Cloudinary!\nID: {pet_id}\nFoto URL: {foto_url}" | |
| def buscar_mascota(imagen_pil: Image.Image, umbral: float = 0.75): | |
| if imagen_pil is None: | |
| return "⚠️ Sube una foto del perro encontrado.", None | |
| vector = get_average_embedding([imagen_pil]) | |
| resultados = qdrant.query_points( | |
| collection_name=COLLECTION_NAME, | |
| query=vector, | |
| limit=3, | |
| score_threshold=umbral | |
| ).points | |
| if not resultados: | |
| return "❌ No se encontraron coincidencias por encima del umbral.", None | |
| texto = "🎉 ¡POSIBLES COINCIDENCIAS ENCONTRADAS!\n\n" | |
| for idx, hit in enumerate(resultados, 1): | |
| similitud = round(hit.score * 100, 2) | |
| texto += f"#{idx} | Coincidencia: {similitud}%\n" | |
| texto += f"🐶 Nombre: {hit.payload.get('nombre')}\n" | |
| texto += f"📞 Contacto: {hit.payload.get('contacto')}\n" | |
| texto += f"📝 Notas: {hit.payload.get('descripcion')}\n\n" | |
| return texto, resultados[0].payload.get("foto_url") | |
| # ==================== 5. INTERFAZ WEB Y API ==================== | |
| with gr.Blocks(title="Buscador IA de Mascotas") as demo: | |
| gr.Markdown("# 🐾 Sistema Híbrido de Reconocimiento de Mascotas") | |
| with gr.Tab("1. Reportar Mascota Perdida"): | |
| in_nombre = gr.Textbox(label="Nombre de la mascota") | |
| in_contacto = gr.Textbox(label="Teléfono / WhatsApp del dueño") | |
| in_desc = gr.Textbox(label="Señas particulares (collar, zona, etc.)") | |
| in_fotos = gr.File(file_count="multiple", file_types=["image"], label="Sube 1 a 3 fotos (ej. cachorro y adulto)") | |
| btn_reg = gr.Button("Guardar Mascota Perdida", variant="primary") | |
| out_reg = gr.Textbox(label="Estado del registro") | |
| btn_reg.click( | |
| fn=registrar_mascota, | |
| inputs=[in_nombre, in_contacto, in_desc, in_fotos], | |
| outputs=out_reg, | |
| api_name="registrar" | |
| ) | |
| with gr.Tab("2. Buscar Perro Encontrado en la Calle"): | |
| in_buscar_img = gr.Image(type="pil", label="Sube la foto tomada en la calle") | |
| in_umbral = gr.Slider(0.60, 0.95, value=0.75, step=0.01, label="Sensibilidad (0.75 recomendado para cambios de edad)") | |
| btn_bus = gr.Button("Buscar Coincidencias", variant="primary") | |
| out_bus_txt = gr.Textbox(label="Resultados", lines=8) | |
| out_bus_img = gr.Image(label="Foto de la mejor coincidencia") | |
| btn_bus.click( | |
| fn=buscar_mascota, | |
| inputs=[in_buscar_img, in_umbral], | |
| outputs=[out_bus_txt, out_bus_img], | |
| api_name="buscar" | |
| ) | |
| demo.launch() |