image-classifier/design_code_classifier.py

"""
design_code_classifier.py — USPTO Design Search Code classifier (RAG version)
==============================================================================
Uses a two-stage RAG flow to classify uploaded trademark images into USPTO
Design Search Codes — with the guarantee that returned codes are real USPTO
codes pulled from the official manual, not invented by the model.

ARCHITECTURE
------------
Stage 1 (visual description):
    Claude vision model looks at the image and produces a structured natural-
    language description of what's in it.

Vector retrieval:
    The description is embedded via Voyage and matched against pre-embedded
    USPTO code descriptions. Top-K candidates are pulled.

Stage 2 (constrained selection):
    Claude sees the image AGAIN alongside the candidate codes (with their
    official descriptions) and selects which apply, with confidence scores
    and rationales.

This means Claude is *picking from a verified menu*, not free-styling. No
hallucinated codes.

DESIGN PHILOSOPHY
-----------------
This module is fully STANDALONE — does not import from or modify
nst_comparison.py. The output is structured JSON; how it gets plugged into
the matching algorithm is a deliberate decision left to the caller.

REQUIREMENTS
------------
    pip install anthropic voyageai pillow numpy python-dotenv

ENV VARS
--------
    ANTHROPIC_API_KEY — Anthropic API key
    VOYAGE_API_KEY    — Voyage AI API key (for query embedding only;
                        codes were pre-embedded by build_code_embeddings.py)
    CLASSIFIER_MODEL  — Optional: override Claude model (default: claude-sonnet-4-6)

PREREQUISITES
-------------
You must run scrape_uspto_design_codes.py and build_code_embeddings.py first,
producing uspto_code_embeddings.pkl in the same directory as this module.

QUICK CLI TEST
--------------
    python design_code_classifier.py path/to/logo.jpg

INTEGRATION (the "subtle change" path for nst_comparison.py)
------------------------------------------------------------
    # ── At top of nst_comparison.py ──
    from .design_code_classifier import classify_image

    # ── Inside execute_model(), after np_img is created ──
    classification = await classify_image(np_img)
    high_conf_codes = classification.high_confidence_codes(threshold=0.7)

    # ── Add design_search_codes to the SELECT, then filter results: ──
    if high_conf_codes:
        response.data = [
            r for r in response.data
            if r.get("design_search_codes") and any(
                c in r["design_search_codes"] for c in high_conf_codes
            )
        ]
"""

import os
import io
import json
import base64
import pickle
import logging
import asyncio
from pathlib import Path
from typing import Optional, Any, List
from dataclasses import dataclass, asdict, field

import numpy as np
from anthropic import AsyncAnthropic
from PIL import Image
from dotenv import load_dotenv

try:
    import voyageai
except ImportError:
    raise ImportError("Voyage AI SDK required. Run: pip install voyageai")

# Optional FastAPI types — only loaded if FastAPI is present
try:
    from fastapi import UploadFile, HTTPException
    HAS_FASTAPI = True
except ImportError:
    HAS_FASTAPI = False
    UploadFile = None  # type: ignore
    HTTPException = None  # type: ignore

# ============================================================================
# CONFIG
# ============================================================================

env_path = Path(__file__).parent / ".env"
load_dotenv(dotenv_path=env_path)

ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
VOYAGE_API_KEY = os.getenv("VOYAGE_API_KEY")

DEFAULT_MODEL = os.getenv("CLASSIFIER_MODEL", "claude-sonnet-4-6")
MAX_TOKENS_STAGE1 = 1024
MAX_TOKENS_STAGE2 = 2048

TOP_K_CANDIDATES = 30
MIN_CONFIDENCE = 0.5

MAX_IMAGE_DIMENSION = 1568

EMBEDDINGS_PATH = Path(__file__).parent / "uspto_code_embeddings.pkl"

logger = logging.getLogger("design_classifier")


# ============================================================================
# CODE INDEX (loads embeddings once, supports vector search)
# ============================================================================

class CodeIndex:
    """In-memory vector index over USPTO design codes.

    Loaded once on first use. All requests share the same index — no per-
    request file I/O.
    """

    _instance: Optional["CodeIndex"] = None

    def __init__(self):
        if not EMBEDDINGS_PATH.exists():
            raise FileNotFoundError(
                f"{EMBEDDINGS_PATH.name} not found. "
                "Run scrape_uspto_design_codes.py + build_code_embeddings.py first."
            )

        with EMBEDDINGS_PATH.open("rb") as f:
            payload = pickle.load(f)

        self.codes: List[str] = payload["codes"]
        self.descriptions: List[str] = payload["descriptions"]
        self.categories: List[str] = payload["categories"]
        self.embeddings: np.ndarray = payload["embeddings"]
        self.metadata: dict = payload["metadata"]

        # Pre-normalize for fast cosine similarity (just dot product after this)
        norms = np.linalg.norm(self.embeddings, axis=1, keepdims=True)
        norms[norms == 0] = 1.0
        self.embeddings_normalized = self.embeddings / norms

        self.code_to_description = dict(zip(self.codes, self.descriptions))

        logger.info(
            f"📚 Loaded code index: {len(self.codes)} codes, "
            f"{self.metadata['dimension']}-dim {self.metadata['model']} embeddings"
        )

    @classmethod
    def get(cls) -> "CodeIndex":
        if cls._instance is None:
            cls._instance = CodeIndex()
        return cls._instance

    def search(self, query_embedding: np.ndarray, top_k: int = TOP_K_CANDIDATES) -> List[dict]:
        """Find top-K most similar codes to the query embedding.

        Returns list of {code, description, category, similarity} dicts,
        sorted descending by similarity.
        """
        q = query_embedding / (np.linalg.norm(query_embedding) + 1e-12)
        sims = self.embeddings_normalized @ q
        top_idx = np.argpartition(sims, -top_k)[-top_k:]
        top_idx = top_idx[np.argsort(-sims[top_idx])]

        return [
            {
                "code": self.codes[i],
                "description": self.descriptions[i],
                "category": self.categories[i],
                "similarity": float(sims[i]),
            }
            for i in top_idx
        ]


# ============================================================================
# IMAGE NORMALIZATION
# ============================================================================

def _normalize_image(image_input: Any) -> tuple[str, str]:
    """Convert various image input types to (base64_data, media_type).

    Accepts: bytes, PIL.Image, numpy.ndarray, file path.
    """
    if hasattr(image_input, "shape") and hasattr(image_input, "dtype"):
        arr = image_input
        if arr.dtype != np.uint8:
            if arr.max() <= 1.0:
                arr = (arr * 255).astype(np.uint8)
            else:
                arr = arr.astype(np.uint8)
        img = Image.fromarray(arr)
    elif isinstance(image_input, bytes):
        img = Image.open(io.BytesIO(image_input))
    elif isinstance(image_input, Image.Image):
        img = image_input
    elif isinstance(image_input, (str, Path)):
        img = Image.open(image_input)
    else:
        raise TypeError(
            f"Unsupported image input type: {type(image_input)}. "
            "Use bytes, PIL.Image, numpy.ndarray, or a file path."
        )

    if img.mode != "RGB":
        img = img.convert("RGB")

    if max(img.size) > MAX_IMAGE_DIMENSION:
        img.thumbnail((MAX_IMAGE_DIMENSION, MAX_IMAGE_DIMENSION), Image.Resampling.LANCZOS)
        logger.debug(f"Resized image to {img.size}")

    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=90)
    base64_data = base64.standard_b64encode(buf.getvalue()).decode("utf-8")
    return base64_data, "image/jpeg"


# ============================================================================
# RESULT TYPES
# ============================================================================

@dataclass
class DesignCode:
    code: str
    description: str
    confidence: float
    rationale: str


@dataclass
class ClassificationResult:
    codes: List[DesignCode] = field(default_factory=list)
    image_description: str = ""
    primary_category: Optional[str] = None
    candidate_codes_considered: List[str] = field(default_factory=list)

    def high_confidence_codes(self, threshold: float = 0.7) -> List[str]:
        """Return just the code strings above the confidence threshold."""
        return [c.code for c in self.codes if c.confidence >= threshold]

    def to_dict(self) -> dict:
        return {
            "codes": [asdict(c) for c in self.codes],
            "image_description": self.image_description,
            "primary_category": self.primary_category,
            "candidate_codes_considered": self.candidate_codes_considered,
        }


# ============================================================================
# STAGE 1 — IMAGE → NATURAL LANGUAGE DESCRIPTION
# ============================================================================

STAGE1_SYSTEM_PROMPT = """You are an expert trademark image analyst. Your job is to describe trademark logos in a way that will help retrieve relevant USPTO Design Search Codes.

When you describe an image, focus on the visual elements a USPTO examiner would code:
- Living things: humans, animals, plants (be specific — "lion's head" not just "animal")
- Geometric shapes: circles, squares, triangles, lines (and their arrangement)
- Objects: tools, vehicles, buildings, food, clothing (be specific)
- Symbols: arrows, crosses, stars, hearts, mathematical symbols
- Text/letters: note their presence and style (stylized, in a frame, etc.)
- Colors: only if the color is a distinctive design element
- Composition: what's framing what, what's stylized vs. realistic

Output a single paragraph (3-6 sentences). Be precise and use concrete vocabulary that a search index would match. Don't editorialize about meaning or branding — just describe what's literally visible.

Output the description as plain text. No JSON, no prose preamble."""


async def stage1_describe_image(
    base64_data: str,
    media_type: str,
    client: AsyncAnthropic,
    model: str,
) -> str:
    """Have Claude describe the image in retrieval-optimized natural language."""
    response = await client.messages.create(
        model=model,
        max_tokens=MAX_TOKENS_STAGE1,
        system=STAGE1_SYSTEM_PROMPT,
        messages=[{
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": media_type,
                        "data": base64_data,
                    },
                },
                {"type": "text", "text": "Describe this trademark image."},
            ],
        }],
    )
    description = response.content[0].text.strip()
    logger.info(f"Stage 1 description ({len(description)} chars): {description[:120]}...")
    return description


# ============================================================================
# RETRIEVAL — VOYAGE EMBED + VECTOR SEARCH
# ============================================================================

def embed_query(text: str, voyage_client) -> np.ndarray:
    """Embed a single query string with Voyage. Returns a 1D numpy array.

    NOTE: model name MUST match what was used in build_code_embeddings.py.
    Reading the metadata from the index ensures consistency.
    """
    index = CodeIndex.get()
    embedding_model = index.metadata["model"]

    result = voyage_client.embed(
        texts=[text],
        model=embedding_model,
        input_type="query",
    )
    return np.array(result.embeddings[0], dtype=np.float32)


# ============================================================================
# STAGE 2 — IMAGE + CANDIDATES → STRUCTURED CODE SELECTION
# ============================================================================

def _stage2_system_prompt(candidates: List[dict]) -> str:
    """Build the system prompt with the candidate codes inline."""
    candidate_lines = "\n".join(
        f"  {c['code']}  —  {c['description']}"
        for c in candidates
    )

    return f"""You are an expert USPTO Trademark Design Search Code classifier.

Below is a list of candidate USPTO Design Search Codes that may apply to the image. Your job is to look at the image and select ONLY the codes that genuinely apply — based on what is actually visible.

CANDIDATE CODES:
{candidate_lines}

CRITICAL RULES:
- ONLY return codes from the candidate list above. Do not invent codes that are not in the list.
- Be precise: include a code only if the corresponding visual element is clearly present.
- Most logos use 2–6 codes total. Don't pad — quality over quantity.
- If the logo contains text or letters, ALWAYS look at category 27 codes in the candidates.
- If the logo has a geometric frame around the elements, ALWAYS look at category 26 codes in the candidates.

Return ONLY valid JSON in this exact structure:
{{
  "codes": [
    {{
      "code": "XX.YY.ZZ",
      "confidence": 0.95,
      "rationale": "Brief explanation of what in the image triggered this code"
    }}
  ],
  "primary_category": "XX"
}}

CONFIDENCE SCORING:
- 0.90–1.00: Element unmistakably present
- 0.70–0.89: Element clearly present, minor interpretation involved
- 0.50–0.69: Element likely present but ambiguous
- Below 0.50: DO NOT include the code

Output valid JSON only. No prose before or after."""


async def stage2_select_codes(
    base64_data: str,
    media_type: str,
    candidates: List[dict],
    client: AsyncAnthropic,
    model: str,
) -> dict:
    """Have Claude pick the applicable codes from the candidate list."""
    response = await client.messages.create(
        model=model,
        max_tokens=MAX_TOKENS_STAGE2,
        system=_stage2_system_prompt(candidates),
        messages=[{
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": media_type,
                        "data": base64_data,
                    },
                },
                {
                    "type": "text",
                    "text": "Select the applicable codes from the candidates. Return JSON only.",
                },
            ],
        }],
    )

    raw_text = response.content[0].text.strip()

    # Strip markdown fences if present
    if raw_text.startswith("```"):
        lines = raw_text.split("\n")
        if lines[-1].startswith("```"):
            lines = lines[1:-1]
        else:
            lines = lines[1:]
        raw_text = "\n".join(lines).strip()

    try:
        return json.loads(raw_text)
    except json.JSONDecodeError as e:
        logger.error(f"Stage 2 JSON parse error: {e}")
        logger.error(f"Raw response: {raw_text[:500]}")
        return {"codes": [], "primary_category": None}


# ============================================================================
# PUBLIC API
# ============================================================================

async def classify_image(
    image_input: Any,
    model: str = DEFAULT_MODEL,
    min_confidence: float = MIN_CONFIDENCE,
    top_k: int = TOP_K_CANDIDATES,
) -> ClassificationResult:
    """Classify a trademark image into USPTO Design Search Codes via two-stage RAG.

    Args:
        image_input:    bytes, PIL.Image, file path, or numpy.ndarray
        model:          Claude model (default: claude-sonnet-4-6)
        min_confidence: Drop codes below this confidence (default: 0.5)
        top_k:          How many candidates to surface to stage 2 (default: 30)

    Returns:
        ClassificationResult with verified codes only — Claude cannot return
        codes that aren't in USPTO's actual vocabulary.
    """
    if not ANTHROPIC_API_KEY:
        raise ValueError("ANTHROPIC_API_KEY environment variable not set")
    if not VOYAGE_API_KEY:
        raise ValueError("VOYAGE_API_KEY environment variable not set")

    # Lazy-load the index (reused across all classifications)
    index = CodeIndex.get()

    # Normalize image once — used in both stages
    base64_data, media_type = _normalize_image(image_input)

    anthropic_client = AsyncAnthropic(api_key=ANTHROPIC_API_KEY)
    voyage_client = voyageai.Client(api_key=VOYAGE_API_KEY)

    # ── Stage 1: describe ──
    description = await stage1_describe_image(
        base64_data, media_type, anthropic_client, model
    )

    # ── Retrieval: embed + vector search ──
    query_emb = embed_query(description, voyage_client)
    candidates = index.search(query_emb, top_k=top_k)
    logger.info(
        f"Retrieved {len(candidates)} candidate codes "
        f"(top similarity: {candidates[0]['similarity']:.3f})"
    )

    # ── Stage 2: select ──
    selection = await stage2_select_codes(
        base64_data, media_type, candidates, anthropic_client, model
    )

    # ── Build result, validating that returned codes are in the candidate set ──
    candidate_codes = {c["code"] for c in candidates}
    result = ClassificationResult(
        image_description=description,
        primary_category=selection.get("primary_category"),
        candidate_codes_considered=[c["code"] for c in candidates],
    )

    for sel in selection.get("codes", []):
        code = sel.get("code", "").strip()
        confidence = float(sel.get("confidence", 0))

        if confidence < min_confidence:
            continue
        if code not in candidate_codes:
            # Defense in depth: even if Claude tries to invent a code, refuse it
            logger.warning(f"Claude returned code {code} not in candidate set — dropping")
            continue

        result.codes.append(DesignCode(
            code=code,
            description=index.code_to_description.get(code, ""),
            confidence=confidence,
            rationale=sel.get("rationale", ""),
        ))

    logger.info(
        f"✅ Classified into {len(result.codes)} verified code(s); "
        f"primary category: {result.primary_category}"
    )
    return result


# ============================================================================
# OPTIONAL HELPERS
# ============================================================================

if HAS_FASTAPI:
    async def classify_uploadfile(file: UploadFile, **kwargs) -> ClassificationResult:
        """Convenience wrapper for FastAPI UploadFile inputs."""
        if file.content_type not in ("image/jpeg", "image/png", "image/gif", "image/webp"):
            raise HTTPException(
                status_code=400,
                detail=f"Unsupported image type: {file.content_type}",
            )
        image_bytes = await file.read()
        return await classify_image(image_bytes, **kwargs)


# ============================================================================
# CLI
# ============================================================================

async def _cli():
    import argparse
    parser = argparse.ArgumentParser(description="Test the USPTO design code classifier")
    parser.add_argument("image_path", help="Path to a trademark image")
    parser.add_argument("--model", default=DEFAULT_MODEL, help="Claude model")
    parser.add_argument("--threshold", type=float, default=MIN_CONFIDENCE)
    parser.add_argument("--top-k", type=int, default=TOP_K_CANDIDATES)
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

    result = await classify_image(
        args.image_path,
        model=args.model,
        min_confidence=args.threshold,
        top_k=args.top_k,
    )

    print(json.dumps(result.to_dict(), indent=2))


if __name__ == "__main__":
    asyncio.run(_cli())