from __future__ import annotations import math import os import numpy as np import torch import torch.nn as nn from .schemas import OutfitCandidate, RecommendationContext class OutfitCompatibilityRanker(nn.Module): """ Transformer ranker. Input: outfit_tokens [B, 6, D] = [CONTEXT, USER, TOP, BOTTOM, SHOES, ACCESSORY] attention_mask [B, 6] Output: logits [B, 1] """ def __init__( self, d_model: int = 512, n_layers: int = 4, n_heads: int = 8, dropout: float = 0.1, ) -> None: super().__init__() self.d_model = d_model self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model)) self.slot_embedding = nn.Embedding(7, d_model) encoder_layer = nn.TransformerEncoderLayer( d_model=d_model, nhead=n_heads, dim_feedforward=d_model * 4, dropout=dropout, batch_first=True, activation="gelu", norm_first=True, ) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers) self.head = nn.Sequential( nn.LayerNorm(d_model), nn.Linear(d_model, d_model // 2), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model // 2, 1), ) nn.init.trunc_normal_(self.cls_token, std=0.02) def forward( self, outfit_tokens: torch.Tensor, attention_mask: torch.Tensor, ) -> torch.Tensor: if outfit_tokens.ndim != 3: raise ValueError("outfit_tokens must have shape [B, S, D]") if attention_mask.ndim != 2: raise ValueError("attention_mask must have shape [B, S]") batch_size, seq_len, _ = outfit_tokens.shape cls = self.cls_token.expand(batch_size, 1, self.d_model) tokens = torch.cat([cls, outfit_tokens], dim=1) slot_ids = torch.arange(seq_len + 1, device=tokens.device).unsqueeze(0) tokens = tokens + self.slot_embedding(slot_ids) cls_mask = torch.ones((batch_size, 1), device=attention_mask.device, dtype=attention_mask.dtype) mask = torch.cat([cls_mask, attention_mask], dim=1) encoded = self.encoder(tokens, src_key_padding_mask=mask == 0) return self.head(encoded[:, 0, :]) class NeuralOutfitScorer: """ Uses a trained transformer checkpoint when available. Otherwise falls back to zero-shot geometric scoring over multimodal embeddings so the endpoint stays usable before fine-tuning. """ def __init__( self, d_model: int = 512, checkpoint_path: str | None = None, device: str | None = None, ) -> None: self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") self.model = OutfitCompatibilityRanker(d_model=d_model).to(self.device).eval() self.checkpoint_path = checkpoint_path or os.getenv("FASHION_RANKER_CHECKPOINT") self.is_trained = False self._load_checkpoint_if_available() def score_candidates( self, candidates: list[OutfitCandidate], context_vector: np.ndarray, user_vector: np.ndarray, context: RecommendationContext, ) -> list[OutfitCandidate]: if not candidates: return [] if self.is_trained: return self._score_with_transformer(candidates, context_vector, user_vector, context) return self._score_zero_shot(candidates, context_vector, user_vector, context) def _score_with_transformer( self, candidates: list[OutfitCandidate], context_vector: np.ndarray, user_vector: np.ndarray, context: RecommendationContext, ) -> list[OutfitCandidate]: token_batch = [] mask_batch = [] for candidate in candidates: vectors = [ context_vector, user_vector, candidate.top.vector, candidate.bottom.vector, candidate.shoes.vector if candidate.shoes is not None else np.zeros_like(context_vector), candidate.accessory.vector if candidate.accessory is not None else np.zeros_like(context_vector), ] mask = [ 1, 1, 1, 1, 1 if candidate.shoes is not None else 0, 1 if candidate.accessory is not None else 0, ] token_batch.append(np.stack(vectors, axis=0)) mask_batch.append(mask) with torch.inference_mode(): logits = self.model( torch.tensor(np.stack(token_batch), dtype=torch.float32, device=self.device), torch.tensor(np.asarray(mask_batch), dtype=torch.long, device=self.device), ).squeeze(-1) probs = torch.sigmoid(logits).detach().cpu().numpy() return self._finalize_candidates(candidates, probs, context_vector, user_vector, context) def _score_zero_shot( self, candidates: list[OutfitCandidate], context_vector: np.ndarray, user_vector: np.ndarray, context: RecommendationContext, ) -> list[OutfitCandidate]: scores = [] for candidate in candidates: item_vectors = [slot_item.vector for slot_item in candidate.slot_items()] outfit_centroid = self._normalize(np.mean(np.stack(item_vectors, axis=0), axis=0)) context_alignment = self._cosine(outfit_centroid, context_vector) user_alignment = self._cosine(outfit_centroid, user_vector) pairwise_cohesion = self._pairwise_cohesion(item_vectors) slot_coverage = math.log1p(len(item_vectors)) / math.log1p(4) score = np.mean( np.asarray( [ self._to_unit_interval(context_alignment), self._to_unit_interval(user_alignment), self._to_unit_interval(pairwise_cohesion), slot_coverage, ], dtype=np.float32, ) ) scores.append(float(np.clip(score, 0.0, 1.0))) return self._finalize_candidates(candidates, scores, context_vector, user_vector, context) def _finalize_candidates( self, candidates: list[OutfitCandidate], probs: list[float] | np.ndarray, context_vector: np.ndarray, user_vector: np.ndarray, context: RecommendationContext, ) -> list[OutfitCandidate]: scored = [] for candidate, prob in zip(candidates, probs, strict=False): candidate.score = round(float(prob) * 100.0, 2) candidate.breakdown = self._build_breakdown(candidate, context_vector, user_vector, context) candidate.reason = self._build_reason(candidate, context) candidate.tip = self._build_tip(candidate) scored.append(candidate) scored.sort(key=lambda item: item.score, reverse=True) return scored def _load_checkpoint_if_available(self) -> None: if not self.checkpoint_path or not os.path.isfile(self.checkpoint_path): return try: payload = torch.load(self.checkpoint_path, map_location=self.device) state_dict = payload.get("model_state_dict", payload) if isinstance(payload, dict) else payload self.model.load_state_dict(state_dict, strict=False) self.model.eval() self.is_trained = True except Exception: self.is_trained = False def _build_breakdown( self, candidate: OutfitCandidate, context_vector: np.ndarray, user_vector: np.ndarray, context: RecommendationContext, ) -> dict[str, float]: item_vectors = [slot_item.vector for slot_item in candidate.slot_items()] outfit_centroid = self._normalize(np.mean(np.stack(item_vectors, axis=0), axis=0)) context_alignment = self._to_score(self._cosine(outfit_centroid, context_vector)) user_affinity = self._to_score(self._cosine(outfit_centroid, user_vector)) visual_cohesion = self._to_score(self._pairwise_cohesion(item_vectors)) top_bottom_compat = self._to_score(self._cosine(candidate.top.vector, candidate.bottom.vector)) occasion_fit = self._occasion_slot_fit(candidate, context.occasion) return { "color": round((visual_cohesion + top_bottom_compat) / 2.0, 2), "style": round((context_alignment + visual_cohesion) / 2.0, 2), "occasion": round((occasion_fit + context_alignment) / 2.0, 2), "context_alignment": round(context_alignment, 2), "user_affinity": round(user_affinity, 2), "visual_cohesion": round(visual_cohesion, 2), } def _build_reason(self, candidate: OutfitCandidate, context: RecommendationContext) -> str: parts = [ f"{candidate.top.item.get('color', 'Unknown')} {candidate.top.item.get('category', 'Topwear')}", f"{candidate.bottom.item.get('color', 'Unknown')} {candidate.bottom.item.get('category', 'Bottomwear')}", ] if candidate.shoes is not None: parts.append( f"{candidate.shoes.item.get('color', 'Unknown')} {candidate.shoes.item.get('category', 'Footwear')}" ) if candidate.accessory is not None: parts.append( f"{candidate.accessory.item.get('color', 'Unknown')} {candidate.accessory.item.get('category', 'Accessory')}" ) return ( f"Learned multimodal embeddings rate {' + '.join(parts)} as a coherent " f"combination for {context.occasion or 'casual'} context." ) @staticmethod def _build_tip(candidate: OutfitCandidate) -> str: if candidate.score >= 85: return "Strong outfit match. Keep accessories minimal so the silhouette stays clean." if candidate.score >= 70: return "Good base outfit. Add one tonal accessory to reinforce the palette." return "This outfit is acceptable, but one slot can be swapped for stronger style alignment." @staticmethod def _occasion_slot_fit(candidate: OutfitCandidate, occasion: str) -> float: occ = str(occasion or "casual").lower() texts = " ".join(slot.metadata_text.lower() for slot in candidate.slot_items()) if occ in texts: return 95.0 if occ in {"formal", "interview", "business", "office"} and any( token in texts for token in ["shirt", "blazer", "trouser", "loafer"] ): return 88.0 if occ in {"party", "festive", "wedding"} and any( token in texts for token in ["silk", "embroidered", "dress", "kurta"] ): return 88.0 if occ in {"sports", "gym", "active"} and any( token in texts for token in ["sneaker", "jogger", "tee", "hoodie"] ): return 85.0 return 72.0 @staticmethod def _pairwise_cohesion(vectors: list[np.ndarray]) -> float: if len(vectors) < 2: return 0.0 scores = [] for left_index in range(len(vectors)): for right_index in range(left_index + 1, len(vectors)): scores.append(NeuralOutfitScorer._cosine(vectors[left_index], vectors[right_index])) return float(np.mean(np.asarray(scores, dtype=np.float32))) @staticmethod def _cosine(left: np.ndarray, right: np.ndarray) -> float: left_vec = NeuralOutfitScorer._normalize(left) right_vec = NeuralOutfitScorer._normalize(right) return float(np.dot(left_vec, right_vec)) @staticmethod def _to_score(value: float) -> float: return round(100.0 * NeuralOutfitScorer._to_unit_interval(value), 2) @staticmethod def _to_unit_interval(value: float) -> float: return float(np.clip((value + 1.0) / 2.0, 0.0, 1.0)) @staticmethod def _normalize(vector: np.ndarray) -> np.ndarray: arr = np.asarray(vector, dtype=np.float32).reshape(-1) norm = float(np.linalg.norm(arr)) if norm < 1e-8: return np.zeros_like(arr, dtype=np.float32) return arr / norm