Spaces:

j2forge
/

cashreader-backend

Sleeping

File size: 7,709 Bytes

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"}