Update image_embedder.py

image_embedder.py

@@ -21,13 +21,17 @@ from __future__ import annotations
 import os
 import sys
 from typing import Optional
-
+import logging
 import numpy as np
 import torch
 import torch.nn.functional as F
 from PIL import Image
 from huggingface_hub import hf_hub_download
 
+
+log = logging.getLogger("image_embedder")
+logging.basicConfig(level=logging.INFO, format="%(levelname)s [%(name)s] %(message)s")
+
 # ---------------------------------------------------------------------------
 # Make the local 'detree' package importable
 # ---------------------------------------------------------------------------
@@ -37,20 +41,29 @@ if _current_dir not in sys.path:
 
 try:
     import clip as _clip_lib
+    log.info("clip package imported successfully.")
 except ImportError:
-    print("Warning: 'clip' package not found. Image embedding will return zeros.")
+    log.error("'clip' package not found — image embedding will return zeros.")
     _clip_lib = None
 
+
 try:
     from detree.model.clip_projector import CLIPProjector
+    log.info("CLIPProjector imported successfully.")
 except ImportError as _e:
-    print(f"Warning: could not import CLIPProjector ({_e}). Image embedding will return zeros.")
+    log.error(f"Could not import CLIPProjector: {_e} — image embedding will return zeros.")
     CLIPProjector = None
 
 # Hugging face
 _BASE_DIR = "MAS-AI-0000/Authentica"
 _PROJECTOR_DIR = os.path.join(_BASE_DIR, "Lib/Models/Image")
 
+log.info(f"[paths] _BASE_DIR      = {_BASE_DIR!r}")
+log.info(f"[paths] _PROJECTOR_DIR = {_PROJECTOR_DIR!r}  exists={os.path.exists(_PROJECTOR_DIR)}")
+if os.path.exists(_PROJECTOR_DIR):
+    log.info(f"[paths] _PROJECTOR_DIR contents: {os.listdir(_PROJECTOR_DIR)}")
+
+
 # ---------------------------------------------------------------------------
 # Configuration
 # ---------------------------------------------------------------------------
@@ -63,6 +76,8 @@ CLIP_PROJECTOR_FILENAME = "Lib/Models/Image/clip_projector.pt"
 # ==== Load assets ====
 clip_projector_path = hf_hub_download(repo_id=REPO_ID, filename=CLIP_PROJECTOR_FILENAME)
 
+log.info(f"[config] device={DEVICE!r}  clip_model={CLIP_MODEL!r}")
+
 # ---------------------------------------------------------------------------
 # Module-level initialisation
 # ---------------------------------------------------------------------------
@@ -75,35 +90,39 @@ _projector:   Optional[object] = None
 def _init() -> None:
     global _clip_model, _clip_prep, _projector
 
+    log.info("_init: starting ImageEmbedder initialisation.")
+
     if _clip_lib is None or CLIPProjector is None:
-        print("ImageEmbedder: required packages unavailable — embedding disabled.")
+        log.error("_init: required packages unavailable — embedding disabled.")
         return
 
     # Load CLIP
+    log.info(f"_init: loading CLIP model {CLIP_MODEL!r} on device={DEVICE!r} ...")
     try:
         _clip_model, _clip_prep = _clip_lib.load(CLIP_MODEL, jit=False, device=DEVICE)
         _clip_model.eval()
         for param in _clip_model.parameters():
             param.requires_grad = False
-        print(f"ImageEmbedder: CLIP ({CLIP_MODEL}) loaded on {DEVICE!r}")
+        log.info(f"_init: CLIP ({CLIP_MODEL}) loaded OK on {DEVICE!r}")
     except Exception as exc:
-        print(f"ImageEmbedder: error loading CLIP: {exc}")
+        log.exception(f"_init: error loading CLIP: {exc}")
         return
 
     # Load CLIPProjector
     if not os.path.exists(_PROJECTOR_DIR):
-        print(f"ImageEmbedder: projector directory not found at {_PROJECTOR_DIR!r}.")
+        log.error(f"_init: projector directory not found at {_PROJECTOR_DIR!r} — embedding disabled.")
         return
 
+    log.info(f"_init: loading CLIPProjector from {_PROJECTOR_DIR!r} ...")
     try:
         _projector = CLIPProjector.from_pretrained(
-            # _PROJECTOR_DIR, device=DEVICE
-            clip_projector_path, device=DEVICE
+            _PROJECTOR_DIR, device=DEVICE
         ).to(DEVICE)
         _projector.eval()
-        print(f"ImageEmbedder: CLIPProjector loaded from {_PROJECTOR_DIR!r}")
+        log.info(f"_init: CLIPProjector loaded OK. "
+                 f"clip_dim={_projector.clip_dim}  target_dim={_projector.target_dim}")
     except Exception as exc:
-        print(f"ImageEmbedder: error loading CLIPProjector: {exc}")
+        log.exception(f"_init: error loading CLIPProjector: {exc}")
 
 
 _init()
@@ -128,18 +147,29 @@ def get_image_embedding(image: Image.Image) -> np.ndarray:
         ``np.ndarray`` of shape ``(1, embedding_dim)`` and dtype float32.
     """
     if _clip_model is None or _projector is None:
+        log.error("get_image_embedding: clip_model or projector is None — returning zeros. Check _init logs.")
         return np.zeros((1, 1), dtype=np.float32)
 
-    image = image.convert("RGB")
-    image_tensor = _clip_prep(image).unsqueeze(0).to(DEVICE)
-
-    # CLIP encode → L2-normalise
-    clip_emb = _clip_model.encode_image(image_tensor).float()
-    clip_emb = F.normalize(clip_emb, dim=-1)
-    clip_emb = clip_emb.float()
-
-    # Project into the DETree embedding space (projector normalises output)
-    projected = _projector(clip_emb, normalize=True)
+    log.info(f"get_image_embedding: input image size={image.size}  mode={image.mode!r}")
+    try:
+        image = image.convert("RGB")
+        image_tensor = _clip_prep(image).unsqueeze(0).to(DEVICE)
+        log.info(f"get_image_embedding: preprocessed tensor shape={tuple(image_tensor.shape)}")
+
+        # CLIP encode → L2-normalise
+        clip_emb = _clip_model.encode_image(image_tensor).float()
+        log.info(f"get_image_embedding: raw CLIP embedding shape={tuple(clip_emb.shape)}  "
+                 f"norm={clip_emb.norm(dim=-1).item():.4f}")
+        clip_emb = F.normalize(clip_emb, dim=-1)
+        clip_emb = clip_emb.float()
+
+        # Project into the DETree embedding space (projector normalises output)
+        projected = _projector(clip_emb, normalize=True)
+        log.info(f"get_image_embedding: projected shape={tuple(projected.shape)}  "
+                 f"norm={projected.norm(dim=-1).item():.4f}")
+    except Exception as exc:
+        log.exception(f"get_image_embedding: failed during inference: {exc}")
+        return np.zeros((1, 1), dtype=np.float32)
 
     return projected.cpu().numpy().astype(np.float32)