Update server1.py

server1.py

@@ -10,7 +10,7 @@ from PIL import Image, UnidentifiedImageError
 import open_clip
 from torchvision import transforms as T
 
-# caches locales
+# caches locales (evitar permisos en /)
 os.environ.setdefault("HF_HOME", "/app/cache")
 os.environ.setdefault("XDG_CACHE_HOME", "/app/cache")
 os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/app/cache/huggingface")
@@ -26,6 +26,7 @@ os.environ["MKL_NUM_THREADS"] = "1"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE  = torch.float16 if DEVICE == "cuda" else torch.float32
 
+# rutas a embeddings
 MODEL_EMB_PATH = os.getenv("MODEL_EMB_PATH", "text_embeddings_modelos_h14.pt")
 VERS_EMB_PATH  = os.getenv("VERS_EMB_PATH",  "text_embeddings_h14.pt")
 
@@ -57,7 +58,7 @@ def _ensure_label_list(x):
 def _load_embeddings(path: str):
     ckpt = torch.load(path, map_location="cpu")
     labels = _ensure_label_list(ckpt["labels"])
-    embeds = ckpt["embeddings"].to("cpu")   # guardados como fp16
+    embeds = ckpt["embeddings"].to("cpu")  # guardados como fp16; los castearemos mas tarde
     embeds = embeds / embeds.norm(dim=-1, keepdim=True)
     return labels, embeds
 
@@ -76,13 +77,14 @@ def _encode_image(img_tensor: torch.Tensor) -> torch.Tensor:
 
 def _predict_top(text_feats_dev: torch.Tensor, text_labels: list[str], image_tensor: torch.Tensor, topk: int = 1):
     img_f = _encode_image(image_tensor)
-    # casteamos embeddings al mismo dtype que la imagen
+    # asegurar mismo device y dtype
     text_feats_dev = text_feats_dev.to(device=img_f.device, dtype=img_f.dtype)
     sim = (100.0 * img_f @ text_feats_dev.T).softmax(dim=-1)[0]
     vals, idxs = torch.topk(sim, k=topk)
-    return [{"label": text_labels[i], "confidence": round(float(v)*100.0, 2)} for v, i in zip(vals, idxs)]
+    return [{"label": text_labels[i], "confidence": round(float(v) * 100.0, 2)} for v, i in zip(vals, idxs)]
 
 def process_image_bytes(image_bytes: bytes):
+    # devuelve solo el dict vehicle: brand/model/version
     if not image_bytes or len(image_bytes) < 128:
         raise UnidentifiedImageError("imagen invalida")
 
@@ -95,7 +97,7 @@ def process_image_bytes(image_bytes: bytes):
     # paso 1: top-1 modelo
     model_feats_dev = model_embeddings.to(device=DEVICE, dtype=DTYPE)
     top_model = _predict_top(model_feats_dev, model_labels, img_tensor, topk=1)[0]
-    modelo_full = top_model["label"]; conf_m = top_model["confidence"]
+    modelo_full = top_model["label"]
 
     partes = modelo_full.split(" ", 1)
     marca  = partes[0] if len(partes) >= 1 else ""
@@ -105,8 +107,9 @@ def process_image_bytes(image_bytes: bytes):
     matches = [(lab, idx) for idx, lab in enumerate(version_labels) if lab.startswith(modelo_full)]
     if not matches:
         return {
-            "marca": marca.upper(), "modelo": modelo.title(),
-            "version": "", "confianza_modelo": conf_m, "confianza_version": 0.0
+            "brand": marca.upper(),
+            "model": modelo.title(),
+            "version": ""
         }
 
     idxs = [i for _, i in matches]
@@ -115,20 +118,20 @@ def process_image_bytes(image_bytes: bytes):
 
     # paso 3: top-1 version
     top_ver = _predict_top(embeds_sub, labels_sub, img_tensor, topk=1)[0]
-    raw = top_ver["label"]; conf_v = top_ver["confidence"]
+    raw = top_ver["label"]
 
     prefix = modelo_full + " "
     ver = raw[len(prefix):] if raw.startswith(prefix) else raw
     ver = ver.split(" ")[0]
-    if conf_v < 25.0:
-        ver = "Version no identificada con suficiente confianza"
+
+    # si baja confianza, no rellenamos version
+    if top_ver["confidence"] < 25.0:
+        ver = ""
 
     return {
-        "marca": marca.upper(),
-        "modelo": modelo.title(),
-        "version": ver.title() if ver else "",
-        "confianza_modelo": conf_m,
-        "confianza_version": conf_v
+        "brand": marca.upper(),
+        "model": modelo.title(),
+        "version": ver.title() if ver else ""
     }
 
 # ============== endpoints ==============
@@ -137,22 +140,34 @@ def root():
     return {"status": "ok", "device": DEVICE}
 
 @app.post("/predict")
-async def predict(front: UploadFile = File(None), back: Optional[UploadFile] = File(None), request: Request = None):
+async def predict(front: UploadFile = File(None),
+                  back: Optional[UploadFile] = File(None),
+                  request: Request = None):
     try:
         if request:
             print("headers:", dict(request.headers))
         if front is None:
-            return JSONResponse(content={"code": 400, "error": "faltan archivos: 'front' es obligatorio"}, status_code=200)
+            return JSONResponse(
+                content={"code": 400, "error": "faltan archivos: 'front' es obligatorio"},
+                status_code=200
+            )
 
         front_bytes = await front.read()
         if back is not None:
             _ = await back.read()
 
-        result = process_image_bytes(front_bytes)
-        return JSONResponse(content={"code": 200, "data": result})
+        vehicle = process_image_bytes(front_bytes)
+        return JSONResponse(
+            content={"code": 200, "data": {"vehicle": vehicle}},
+            status_code=200
+        )
 
     except Exception as e:
         print("EXCEPTION:", repr(e))
         traceback.print_exc()
-        return JSONResponse(content={"code": 404, "data": {}, "error": str(e)}, status_code=200)
+        return JSONResponse(
+            content={"code": 404, "data": {}, "error": str(e)},
+            status_code=200
+        )
+