Spaces:
Sleeping
Sleeping
| import io | |
| import os | |
| import boto3 | |
| import gc | |
| import sys | |
| import hashlib | |
| from datetime import datetime, timezone, timedelta | |
| from PIL import Image | |
| from botocore.config import Config | |
| from ultralytics import YOLO | |
| from fastapi import FastAPI, UploadFile | |
| from fastapi.middleware.cors import CORSMiddleware | |
| # ========================================== | |
| # 1. CONFIGURACIÓN DE CLOUDFLARE R2 | |
| # ========================================== | |
| R2_BUCKET_NAME = os.getenv("R2_BUCKET_NAME") | |
| R2_ENDPOINT_URL = os.getenv("R2_ENDPOINT_URL") | |
| R2_ACCESS_KEY = os.getenv("R2_ACCESS_KEY") | |
| R2_SECRET_KEY = os.getenv("R2_SECRET_KEY") | |
| model_configs = { | |
| "VEF": { | |
| "r2_key": "files/md5/39/398b1e461a92a06981e0a1127a9f56", | |
| "path": "/home/user/models/train/VEF_model_13f/weights/best.onnx" | |
| }, | |
| "USD": { | |
| "r2_key": "files/md5/8b/50f8fa0b30f38111e8f77fcb396639", | |
| "path": "/home/user/models/train/USD_model_plus_01/weights/best.onnx" | |
| }, | |
| "INFERENCIA": { | |
| "r2_key": "files/md5/1d/ee121124dc76d37c3a5efc4993f961", | |
| "path": "/home/user/models/train/USD_VEF_Model_01j/weights/best.onnx" | |
| } | |
| } | |
| # ========================================== | |
| # 2. FUNCIONES AUXILIARES | |
| # ========================================== | |
| def get_s3_client(): | |
| """Retorna un cliente s3 configurado para Cloudflare R2""" | |
| return boto3.client( | |
| service_name="s3", | |
| endpoint_url=R2_ENDPOINT_URL, | |
| aws_access_key_id=R2_ACCESS_KEY, | |
| aws_secret_access_key=R2_SECRET_KEY, | |
| config=Config(signature_version="s3v4") | |
| ) | |
| # ========================================== | |
| # 3. DESCARGA CONSTRUCTORA Y DESTRUCTIVA | |
| # ========================================== | |
| def download_models_from_r2(): | |
| s3_client = get_s3_client() | |
| for name, config in model_configs.items(): | |
| os.makedirs(os.path.dirname(config["path"]), exist_ok=True) | |
| if not os.path.exists(config["path"]): | |
| print(f"Descargando modelo {name} desde Cloudflare R2...") | |
| try: | |
| s3_client.download_file( | |
| Bucket=R2_BUCKET_NAME, | |
| Key=config["r2_key"], | |
| Filename=config["path"] | |
| ) | |
| print(f"¡Modelo {name} descargado con éxito!") | |
| except Exception as e: | |
| print(f"❌ Error al descargar el modelo {name}: {e}") | |
| raise e | |
| else: | |
| print(f"El modelo {name} ya existe localmente. Cargando...") | |
| del s3_client | |
| gc.collect() | |
| download_models_from_r2() | |
| # ========================================== | |
| # 4. CARGA DE MODELOS | |
| # ========================================== | |
| models = { | |
| "USD": YOLO(model_configs["USD"]["path"], task="detect"), | |
| "VEF": YOLO(model_configs["VEF"]["path"], task="detect"), | |
| "INFERENCIA": YOLO(model_configs["INFERENCIA"]["path"], task="detect"), | |
| } | |
| classes = { | |
| "USD": [ | |
| "fifty-back", "fifty-front", | |
| "five-back", "five-front", | |
| "one-back", "one-front", | |
| "one_hundred-back", "one_hundred-front", | |
| "ten-back", "ten-front", | |
| "twenty-back", "twenty-front", | |
| ], | |
| "VEF": [ | |
| "five-back-vef", "five-front-vef", | |
| "fifty-back-vef", "fifty-front-vef", | |
| "five_hundred-back-vef", "five_hundred-front-vef", | |
| "one_hundred-back-vef", "one_hundred-front-vef", | |
| "ten-back-vef", "ten-front-vef", | |
| "twenty-back-vef", "twenty-front-vef", | |
| "two_hundred-back-vef", "two_hundred-front-vef", | |
| ], | |
| "INFERENCIA": [ | |
| "dollar_back", "dollar_front", | |
| "vef_back", "vef_front", | |
| ], | |
| } | |
| app = FastAPI() | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| def detection_vef(image: UploadFile): | |
| print("FOTO RECIBIDA") | |
| imageBytes = image.file.read() | |
| imageStream = io.BytesIO(imageBytes) | |
| imageFile = Image.open(imageStream).convert("RGB") | |
| debug_buffer = io.BytesIO() | |
| imageFile.save(debug_buffer, format="JPEG") | |
| which_currency = models["INFERENCIA"].predict(imageFile, verbose=False, imgsz=320, conf=0.10) | |
| if len(which_currency[0].boxes) == 0: | |
| return {"message": "No objects detected"} | |
| currency_label = classes["INFERENCIA"][int(which_currency[0].boxes[0].cls.item())] | |
| if "vef" in currency_label: | |
| currency = "VEF" | |
| else: | |
| currency = "USD" | |
| results = models[currency].predict(imageFile, verbose=False, imgsz=320, conf=0.25) | |
| if len(results[0].boxes) > 0: | |
| # Tomamos el label y la confianza de la primera detección (la de mayor confianza usualmente) | |
| primary_label = classes[currency][int(results[0].boxes[0].cls.item())] | |
| primary_conf = results[0].boxes[0].conf.item() # <--- Extraemos la confianza | |
| # ----------------------------------------------------- | |
| # LÓGICA DE GUARDADO EN R2 (LOGS) | |
| # ----------------------------------------------------- | |
| try: | |
| # Zona horaria de Venezuela (UTC-4) | |
| vzla_tz = timezone(timedelta(hours=-4)) | |
| now = datetime.now(vzla_tz) | |
| # Formato base del nombre | |
| base_filename = f"{now.strftime('%d-%m-%Y_%H%M%S')}_{primary_label}_{primary_conf:.4f}" | |
| # Keys correctas con sus extensiones para R2 | |
| r2_image_key = f"logs/{base_filename}.jpg" | |
| r2_text_key = f"logs/{base_filename}.txt" | |
| # 1. Preparar las coordenadas del TXT en formato YOLO (norm_x_center, norm_y_center, norm_width, norm_height) | |
| log_lines = [] | |
| for box in results[0].boxes: | |
| class_id = int(box.cls.item()) | |
| # xywhn devuelve valores normalizados entre 0 y 1 | |
| x_center, y_center, width, height = box.xywhn[0].tolist() | |
| line = f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n" | |
| log_lines.append(line) | |
| # Convertimos el string acumulado en bytes listos para S3/R2 | |
| text_bytes = "".join(log_lines).encode("utf-8") | |
| # 2. Subir ambos archivos usando el mismo cliente S3 | |
| s3_client = get_s3_client() | |
| # Subir Imagen | |
| s3_client.put_object( | |
| Bucket=R2_BUCKET_NAME, | |
| Key=r2_image_key, | |
| Body=imageBytes, | |
| ContentType="image/jpeg" | |
| ) | |
| print(f"[LOG] Imagen respaldada en R2: {r2_image_key}") | |
| # Subir TXT de anotaciones | |
| s3_client.put_object( | |
| Bucket=R2_BUCKET_NAME, | |
| Key=r2_text_key, | |
| Body=text_bytes, | |
| ContentType="text/plain" # ContentType correcto para archivos de texto limpio | |
| ) | |
| print(f"[LOG] TXT de anotaciones respaldado en R2: {r2_text_key}") | |
| except Exception as e: | |
| print(f"❌ [ERROR] Fallo al subir logs a R2: {e}") | |
| # ----------------------------------------------------- | |
| boxes = [ | |
| { | |
| "label": classes[currency][int(box.cls.item())], | |
| "confidence": box.conf.item(), | |
| "bbox": box.xyxy.tolist() | |
| } | |
| for box in results[0].boxes | |
| ] | |
| print(boxes) | |
| return {"detections": boxes} | |
| else: | |
| return {"message": "No objects detected"} | |
| def status(): | |
| print("200 OK") | |
| return {"message": "200 OK"} |