""" 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())