AtharvaRJ/few-shot-object-identification

1"""2prototype_classifier.py3 4Implements a Prototypical-Networks-style few-shot classifier on top of5CLIP embeddings.6 7Core idea:8  - Each class is represented by a "prototype" vector: the mean of its9    few-shot example image embeddings, optionally fused with the text10    description embedding.11  - A new image is classified by nearest prototype (cosine similarity).12  - If the best similarity is below `confidence_threshold`, the image is13    flagged as "unknown" rather than forced into the nearest class.14"""15 16from __future__ import annotations17from dataclasses import dataclass, field18from datetime import datetime, timezone19import numpy as np20 21 22@dataclass23class ClassPrototype:24    name: str25    prototype: np.ndarray            # (D,) L2-normalized26    example_embeddings: np.ndarray   # (N, D) the raw examples that built it27    text_embedding: np.ndarray | None = None28    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())29    n_examples: int = 030 31    def __post_init__(self):32        self.n_examples = len(self.example_embeddings)33 34 35@dataclass36class PredictionResult:37    predicted_class: str | None38    confidence: float39    all_scores: dict[str, float]40    is_unknown: bool41 42 43class PrototypeClassifier:44    """45    image_text_weight: how much the text-description embedding contributes46    to the final prototype vector, vs. the image examples themselves.47    0.0 = image-only prototype, 1.0 = text-only. 0.3 is a reasonable default48    since the text description helps disambiguate visually-similar classes49    (e.g. "wine bottle" vs "water bottle") that pure image averaging blurs.50    """51 52    def __init__(self, confidence_threshold: float = 0.22, image_text_weight: float = 0.3):53        self.confidence_threshold = confidence_threshold54        self.image_text_weight = image_text_weight55        self.classes: dict[str, ClassPrototype] = {}56 57    def add_class(58        self,59        name: str,60        example_embeddings: np.ndarray,61        text_embedding: np.ndarray | None = None,62    ) -> ClassPrototype:63        """Register a new class from 1-3 (or more) example image embeddings."""64        if name in self.classes:65            raise ValueError(f"Class '{name}' already exists. Use update_class instead.")66 67        image_proto = example_embeddings.mean(axis=0)68        image_proto = image_proto / np.linalg.norm(image_proto)69 70        if text_embedding is not None:71            w = self.image_text_weight72            fused = (1 - w) * image_proto + w * text_embedding73            fused = fused / np.linalg.norm(fused)74        else:75            fused = image_proto76 77        proto = ClassPrototype(78            name=name,79            prototype=fused,80            example_embeddings=example_embeddings,81            text_embedding=text_embedding,82        )83        self.classes[name] = proto84        return proto85 86    def predict(self, query_embedding: np.ndarray) -> PredictionResult:87        """Classify a single query embedding against all stored prototypes."""88        if not self.classes:89            raise RuntimeError("No classes registered yet. Call add_class first.")90 91        scores = {92            name: float(np.dot(query_embedding, proto.prototype))93            for name, proto in self.classes.items()94        }95        best_name = max(scores, key=scores.get)96        best_score = scores[best_name]97        is_unknown = best_score < self.confidence_threshold98 99        return PredictionResult(100            predicted_class=None if is_unknown else best_name,101            confidence=best_score,102            all_scores=scores,103            is_unknown=is_unknown,104        )105 106    def class_names(self) -> list[str]:107        return list(self.classes.keys())