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=["*"], ) @app.post("/detection") 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"} @app.get("/") def status(): print("200 OK") return {"message": "200 OK"}