{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":129543,"databundleVersionId":15525987,"isSourceIdPinned":false},{"sourceType":"kernelVersion","sourceId":306884780,"isSourceIdPinned":false}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"5ab494b7","cell_type":"markdown","source":"# 🐆 Jaguar Re-Identification — Inference & Submission Pipeline\n\n> **Goal:** Use trained EVA-02 models to compute similarity scores between \n> test images and produce a competition submission file.\n\n---\n\n## 📋 Table of Contents\n\n| #  | Section | Description |\n|----|---------|-------------|\n| 1  | Environment Setup & Imports | Importing required dependencies |\n| 2  | Configuration | Inference parameters and checkpoint definitions |\n| 3  | Model Architecture | EVA-02 inference model (no ArcFace) |\n| 4  | Post-Processing | AQE + K-Reciprocal Jaccard Re-Ranking |\n| 5  | Dataset & Feature Extraction | Test data loading and TTA-based embedding extraction |\n| 6  | Main Inference Pipeline | Ensemble → AQE → Re-Ranking → Submission |\n\n---\n\n### Pipeline Flow\n```\n5 Checkpoints (weighted) ──→ Feature-Level Ensemble\n                                      ↓\n                              L2 Normalization\n                                      ↓\n                         Adaptive Query Expansion (AQE)\n                                      ↓\n                      K-Reciprocal Jaccard Re-Ranking\n                                      ↓\n                          Gamma Calibration → CSV\n```\n","metadata":{}},{"id":"d480f8c7","cell_type":"markdown","source":"## 1. Environment Setup & Imports\n\nThe inference pipeline requires fewer libraries compared to training — \nno optimizer, scheduler, or loss functions are needed.\n","metadata":{}},{"id":"fd7595c1","cell_type":"code","source":"# =============================================================================\n# Standard Python libraries\n# =============================================================================\nimport gc                          # Memory management — cleanup after each model\nimport logging                     # Structured log output\nfrom dataclasses import dataclass, field  # Configuration class\nfrom pathlib import Path           # Platform-independent file paths\nfrom typing import List, Optional, Union  # Type hints\n\n# =============================================================================\n# Scientific computing\n# =============================================================================\nimport numpy as np                 # Matrix operations (similarity computation)\nimport pandas as pd                # CSV reading/writing\n\n# =============================================================================\n# Deep learning — PyTorch ecosystem\n# =============================================================================\nimport torch                       # Core deep learning framework\nimport torch.nn as nn              # Neural network layers\nimport torch.nn.functional as F    # Functional API (normalize etc.)\nimport timm                        # Pretrained Vision Transformer models\nfrom torch.utils.data import DataLoader, Dataset  # Data loading infrastructure\n\n# =============================================================================\n# Image processing & Augmentation\n# =============================================================================\nimport albumentations as A         # Image transforms (Resize, Normalize)\nimport cv2                         # Horizontal flip (TTA)\nfrom PIL import Image              # Reading image files\n\n# =============================================================================\n# Utilities\n# =============================================================================\nfrom tqdm.auto import tqdm         # Progress bar\n\n# =============================================================================\n# Logger configuration\n# =============================================================================\nlogging.basicConfig(\n    level=logging.INFO, \n    format='%(asctime)s — %(name)s — %(message)s'\n)\nlogger = logging.getLogger(__name__)\n\nlogger.info(f\"PyTorch: {torch.__version__}\")\nlogger.info(f\"CUDA: {torch.cuda.is_available()}\")\nif torch.cuda.is_available():\n    logger.info(f\"GPU: {torch.cuda.get_device_name(0)}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-27T18:38:19.191061Z","iopub.execute_input":"2026-03-27T18:38:19.191545Z","iopub.status.idle":"2026-03-27T18:38:34.827497Z","shell.execute_reply.started":"2026-03-27T18:38:19.191519Z","shell.execute_reply":"2026-03-27T18:38:34.826656Z"}},"outputs":[],"execution_count":null},{"id":"7794474c","cell_type":"markdown","source":"## 2. Configuration\n\nAll inference pipeline settings are centralized here.\n\n### Ensemble Strategy\nEmbeddings from 5 different checkpoints are combined via **weighted averaging**.  \nThe final epoch (20) gets the highest weight since the LR is at its lowest point in the cosine schedule.\n\n### Post-Processing Chain\n| Technique | Parameter | Description |\n|-----------|-----------|-------------|\n| AQE | threshold=0.70, max_k=7 | Query enrichment with high-similarity neighbors |\n| Re-Ranking | k1=17, weight=0.20 | Re-ranking with Jaccard neighborhood similarity |\n| Calibration | gamma=1.5 | Sharpening the probability distribution |\n","metadata":{}},{"id":"91b9d22e","cell_type":"code","source":"@dataclass\nclass Config:\n    \"\"\"\n    Configuration class holding all inference pipeline settings.\n    \n    Unlike the training Config, this class has no optimizer/scheduler parameters.\n    Instead, it defines ensemble weights and post-processing parameters.\n    \"\"\"\n    \n    # -------------------------------------------------------------------------\n    # File Paths\n    # -------------------------------------------------------------------------\n    # Directory containing trained model weights\n    weights_dir: Path = field(\n        default_factory=lambda: Path(\"/kaggle/input/notebooks/ozanbozyel/jaguar-re-id-eva-02-training\")\n    )\n    # Competition dataset root directory\n    input_root: Path = field(\n        default_factory=lambda: Path(\"/kaggle/input/competitions/round-2-jaguar-reidentification-challenge\")\n    )\n    # Output directory for submission files\n    output_dir: Path = field(default_factory=lambda: Path(\".\"))\n    \n    # -------------------------------------------------------------------------\n    # Model Settings\n    # -------------------------------------------------------------------------\n    model_name: str = \"eva02_large_patch14_448.mim_m38m_ft_in22k_in1k\"\n    img_size: int = 448       # Same resolution as training\n    patch_size: int = 14      # ViT patch size\n    batch_size: int = 4       # Inference batch size\n    num_workers: int = 2      # Data loading worker threads\n    \n    # -------------------------------------------------------------------------\n    # Adaptive Query Expansion (AQE) Parameters\n    # -------------------------------------------------------------------------\n    use_aqe: bool = True           # Enable/disable AQE\n    aqe_threshold: float = 0.70    # Minimum similarity to qualify as neighbor\n    aqe_max_k: int = 7             # Maximum neighbors for query expansion\n    \n    # -------------------------------------------------------------------------\n    # K-Reciprocal Re-Ranking Parameters\n    # -------------------------------------------------------------------------\n    use_re_ranking: bool = True    # Enable/disable Re-Ranking\n    re_rank_k1: int = 17           # Number of neighbors to examine per sample (15-20 ideal)\n    re_rank_weight: float = 0.20   # Jaccard score contribution to final result (20%)\n    \n    # -------------------------------------------------------------------------\n    # Ensemble: Checkpoints and Weights\n    # -------------------------------------------------------------------------\n    # Weighted average of 5 checkpoints for feature-level ensemble\n    # Final epoch (20) gets highest weight — most stable at cosine schedule end\n    checkpoints: List[str] = field(default_factory=lambda: [\n        \"model_eva02_epoch_20.pth\",   # Latest epoch → highest weight\n        \"model_eva02_epoch_17.pth\",   # Earlier epoch — for diversity\n        \"model_eva02_best.pth\",       # Best mAP checkpoint\n        \"model_eva02_epoch_19.pth\",   # Near-final epoch\n        \"model_eva02_epoch_18.pth\",   # Near-final epoch\n    ])\n    weights: List[float] = field(default_factory=lambda: [\n        0.35, 0.25, 0.15, 0.15, 0.10\n    ])\n    \n    # Probability calibration: gamma > 1 → amplifies high scores further\n    calibration_gamma: float = 1.5\n    \n    # -------------------------------------------------------------------------\n    # Computed Properties\n    # -------------------------------------------------------------------------\n    @property\n    def grid_size(self) -> int:\n        \"\"\"ViT spatial grid size: 448 / 14 = 32.\"\"\"\n        return self.img_size // self.patch_size\n    \n    @property\n    def test_dir(self) -> Path:\n        \"\"\"Directory containing test images.\"\"\"\n        return self.input_root / \"test\" if self.input_root.exists() else Path(\"test\")\n    \n    @property\n    def csv_path(self) -> Path:\n        \"\"\"CSV file defining test pairs.\"\"\"\n        return self.input_root / \"test.csv\" if self.input_root.exists() else Path(\"test.csv\")\n    \n    @property\n    def device(self) -> torch.device:\n        \"\"\"Compute device: CUDA if GPU available, otherwise CPU.\"\"\"\n        return torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n\nconfig = Config()\nlogger.info(f\"Model: {config.model_name}\")\nlogger.info(f\"Checkpoints: {len(config.checkpoints)}\")\nlogger.info(f\"Device: {config.device}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-27T18:38:34.828923Z","iopub.execute_input":"2026-03-27T18:38:34.829407Z","iopub.status.idle":"2026-03-27T18:38:34.843416Z","shell.execute_reply.started":"2026-03-27T18:38:34.829380Z","shell.execute_reply":"2026-03-27T18:38:34.842854Z"}},"outputs":[],"execution_count":null},{"id":"77d36145","cell_type":"markdown","source":"## 3. Model Architecture (Inference Version)\n\nKey differences from the training model:\n- **No ArcFace head** — classification is not needed at inference, only embeddings\n- **`pretrained=False`** — weights are loaded from checkpoint, not from ImageNet\n- **No gradient computation** — runs under `torch.no_grad()`\n\n```\nImage [B, 3, 448, 448]\n    ↓\nEVA-02 Backbone → Patch Tokens [B, 1024+1, 1024]\n    ↓\nRemove CLS Token → [B, 1024, 32, 32]\n    ↓\nGeM Pooling → [B, 1024]\n    ↓\nBatchNorm → Normalized Embedding [B, 1024]\n```\n","metadata":{}},{"id":"45b39964","cell_type":"code","source":"class GeM(nn.Module):\n    \"\"\"\n    Generalized Mean Pooling — learnable pooling layer.\n    \n    The p parameter is learned during training and remains fixed at inference.\n    When p > 1, higher activations receive more weight, enabling the model\n    to focus on discriminative regions.\n    \"\"\"\n    def __init__(self, p: float = 3.0, eps: float = 1e-6):\n        super().__init__()\n        self.p = nn.Parameter(torch.ones(1) * p)\n        self.eps = eps\n        \n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        return F.avg_pool2d(\n            x.clamp(min=self.eps).pow(self.p), \n            (x.size(-2), x.size(-1))\n        ).pow(1.0 / self.p)\n\n\nclass EVAReIDModelInference(nn.Module):\n    \"\"\"\n    EVA-02 inference model — produces embeddings only (no ArcFace head).\n    \n    Uses only the backbone + pool + bn layers from the training model\n    (EVAReIDModel). ArcFace weights are automatically filtered during\n    checkpoint loading.\n    \n    Args:\n        config: Inference configuration.\n    \"\"\"\n    def __init__(self, config: Config):\n        super().__init__()\n        self.config = config\n        \n        # Backbone: pretrained=False → weights will come from checkpoint\n        self.backbone = timm.create_model(\n            config.model_name, \n            pretrained=False,        # Loaded from checkpoint\n            num_classes=0,           # No classification layer\n            global_pool=\"\",          # Using our own pooling\n            img_size=config.img_size\n        )\n        self.in_features = self.backbone.num_features  # EVA-02 Large → 1024\n        \n        # Pooling and normalization (same structure as training)\n        self.pool = GeM(p=3)\n        self.bn = nn.BatchNorm1d(self.in_features)\n        nn.init.constant_(self.bn.weight, 1.0)\n        nn.init.constant_(self.bn.bias, 0.0)\n\n    def forward_features(self, x: torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        Extracts spatial feature map from backbone.\n        \n        Removes CLS token (index 0) and reshapes patch tokens to 2D grid.\n        \"\"\"\n        features = self.backbone.forward_features(x)\n        \n        # Remove CLS token (index 0)\n        features = features[:, 1:, :]\n        \n        # [B, N, C] → [B, C, H, W]\n        return features.permute(0, 2, 1).reshape(\n            x.size(0), self.in_features, \n            self.config.grid_size, self.config.grid_size\n        )\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        Forward pass: Image → Normalized embedding vector.\n        \n        Args:\n            x: Input images [B, 3, 448, 448].\n            \n        Returns:\n            Post-BatchNorm embedding [B, 1024].\n        \"\"\"\n        features = self.forward_features(x)\n        pooled = self.pool(features).flatten(start_dim=1)\n        return self.bn(pooled)\n\n\ndef load_checkpoint(model: nn.Module, path: Path) -> Optional[nn.Module]:\n    \"\"\"\n    Loads a training checkpoint into the inference model.\n    \n    Steps:\n    1. Load checkpoint to CPU (to preserve GPU memory)\n    2. Strip 'module.' prefix (DataParallel compatibility)\n    3. Filter out ArcFace layer weights (not needed for inference)\n    4. Load remaining weights with strict=False\n    \n    Args:\n        model: Model to load weights into.\n        path: Checkpoint file path (.pth).\n        \n    Returns:\n        Model with loaded weights, or None if file not found.\n    \"\"\"\n    if not path.exists():\n        logger.warning(f\"Checkpoint not found: {path}\")\n        return None\n    \n    # Load to CPU first → move to GPU later (memory optimization)\n    state_dict = torch.load(path, map_location=\"cpu\", weights_only=True)\n    \n    # Strip DataParallel 'module.' prefix + filter ArcFace weights\n    cleaned_dict = {\n        k[7:] if k.startswith(\"module.\") else k: v \n        for k, v in state_dict.items() \n        if \"arcface\" not in k  # ArcFace head not used in inference\n    }\n    \n    # strict=False: Missing ArcFace keys won't raise errors\n    model.load_state_dict(cleaned_dict, strict=False)\n    logger.info(f\"Checkpoint loaded: {path.name}\")\n    return model\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-27T18:38:34.844286Z","iopub.execute_input":"2026-03-27T18:38:34.844522Z","iopub.status.idle":"2026-03-27T18:38:34.858244Z","shell.execute_reply.started":"2026-03-27T18:38:34.844501Z","shell.execute_reply":"2026-03-27T18:38:34.857538Z"}},"outputs":[],"execution_count":null},{"id":"3fd4ff19","cell_type":"markdown","source":"## 4. Post-Processing: AQE & K-Reciprocal Re-Ranking\n\nTwo techniques for improving raw embeddings after inference:\n\n### 🔹 Adaptive Query Expansion (AQE)\nEnriches each query embedding by averaging it with its high-similarity neighbors.\nThis reduces noise and decreases intra-class variance.\n\n```\nQuery q → Find neighbors with similarity > 0.70 → Average them → New q'\n```\n\n### 🔹 K-Reciprocal Jaccard Re-Ranking\nLooks beyond direct pairwise distance to examine **overlap in neighborhood structures**.\nTwo samples that share the same neighbors are likely from the same class.\n\n```\nJaccard(A, B) = |Neighbors(A) ∩ Neighbors(B)| / |Neighbors(A) ∪ Neighbors(B)|\nFinal score = (1 - λ) × Cosine + λ × Jaccard\n```\n\n> This technique is based on \"Zhong et al., 2017\" and is widely used\n> in Re-ID competitions.\n","metadata":{}},{"id":"969b5082","cell_type":"code","source":"class PostProcessor:\n    \"\"\"\n    Post-inference embedding refinement methods.\n    \n    All methods are defined as static — they carry no state.\n    \"\"\"\n    \n    @staticmethod\n    def adaptive_query_expansion(\n        features: np.ndarray, \n        threshold: float = 0.85, \n        max_k: int = 4\n    ) -> np.ndarray:\n        \"\"\"\n        Adaptive Query Expansion (AQE) — enriches query embeddings.\n        \n        For each embedding:\n        1. Find neighbors with cosine similarity above threshold\n        2. Keep at most max_k neighbors\n        3. Use the average of these neighbors as the new embedding\n        \n        This smooths noisy samples and produces embeddings closer\n        to class centers.\n        \n        Args:\n            features: L2-normalized embedding matrix [N, D].\n            threshold: Minimum similarity for neighbor qualification.\n            max_k: Maximum neighbors per query.\n            \n        Returns:\n            Enriched and re-normalized embeddings [N, D].\n        \"\"\"\n        # L2 normalize (safety check)\n        features = features / (np.linalg.norm(features, axis=1, keepdims=True) + 1e-12)\n        \n        # Full cosine similarity matrix [N, N]\n        sims = np.dot(features, features.T)\n        expanded = np.zeros_like(features)\n        \n        for i in range(len(features)):\n            # Find neighbors above threshold\n            valid_indices = np.where(sims[i] >= threshold)[0]\n            \n            # If too many neighbors, keep only the top max_k\n            if len(valid_indices) > max_k:\n                best_indices = np.argsort(-sims[i])[:max_k]\n                valid_indices = best_indices\n            \n            # Average of neighbors → new embedding\n            expanded[i] = np.mean(features[valid_indices], axis=0)\n        \n        # Re-normalize to unit vectors\n        expanded = expanded / (np.linalg.norm(expanded, axis=1, keepdims=True) + 1e-12)\n        return expanded\n\n    @staticmethod\n    def k_reciprocal_jaccard_similarity(\n        features: np.ndarray, \n        k1: int = 15, \n        jaccard_weight: float = 0.3\n    ) -> np.ndarray:\n        \"\"\"\n        K-Reciprocal Jaccard Re-Ranking — refines similarity using neighborhood structure.\n        \n        Beyond standard cosine similarity, this method also computes the Jaccard\n        similarity of each pair's neighborhood sets and blends both scores.\n        \n        Intuition: If A and B are close AND share the same neighbors,\n        they are much more likely to be the same identity.\n        \n        Formula:\n            final_sim(i,j) = (1 - λ) × cos_sim(i,j) + λ × jaccard(i,j)\n        \n        Args:\n            features: L2-normalized embedding matrix [N, D].\n            k1: Number of neighbors to consider per sample.\n            jaccard_weight: Jaccard score contribution ratio (λ).\n            \n        Returns:\n            Re-ranked similarity matrix [N, N].\n        \"\"\"\n        # Original cosine similarity matrix\n        sim_mat = np.dot(features, features.T)\n        \n        # Find top-k1 neighbor indices for each sample\n        indices = np.argsort(-sim_mat, axis=1)[:, :k1]\n        \n        N = features.shape[0]\n        final_sim_mat = sim_mat.copy()\n        \n        # Compute Jaccard similarity for each pair\n        # Note: Symmetric matrix — only compute upper triangle, then mirror\n        for i in tqdm(range(N), desc=\"K-Reciprocal Re-Ranking\"):\n            k_reciprocal_i = set(indices[i])  # Neighbor set of i\n            \n            for j in range(i + 1, N):  # Symmetry: upper triangle only\n                k_reciprocal_j = set(indices[j])  # Neighbor set of j\n                \n                # Jaccard similarity: Intersection / Union\n                intersection = len(k_reciprocal_i.intersection(k_reciprocal_j))\n                union = len(k_reciprocal_i.union(k_reciprocal_j))\n                j_score = intersection / union if union > 0 else 0.0\n                \n                # Weighted blend of cosine and Jaccard scores\n                final_score = (1 - jaccard_weight) * sim_mat[i, j] + jaccard_weight * j_score\n                \n                # Assign symmetrically\n                final_sim_mat[i, j] = final_score\n                final_sim_mat[j, i] = final_score\n                \n        return final_sim_mat\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-27T18:38:34.860002Z","iopub.execute_input":"2026-03-27T18:38:34.860260Z","iopub.status.idle":"2026-03-27T18:38:34.875253Z","shell.execute_reply.started":"2026-03-27T18:38:34.860238Z","shell.execute_reply":"2026-03-27T18:38:34.874324Z"}},"outputs":[],"execution_count":null},{"id":"6b82a067","cell_type":"markdown","source":"## 5. Dataset & Feature Extraction\n\n### Test-Time Augmentation (TTA)\nFor each image, **original + horizontal flip** versions are produced (2 total).\nThe embeddings of both versions are averaged, yielding more robust representations.\n\n```\nImage → [Original, Flip] → Model → [emb_orig, emb_flip] → Average → Final Embedding\n```\n\n### Mixed Precision Inference\nWhen GPU is available, `torch.amp.autocast('cuda')` enables FP16 inference — \nfaster execution with no impact on embedding quality.\n","metadata":{}},{"id":"ca13a818","cell_type":"code","source":"class InferenceDataset(Dataset):\n    \"\"\"\n    Test image Dataset — with TTA (horizontal flip) support.\n    \n    Each __getitem__ call returns 2 tensors:\n    - Original image (resize + normalize)\n    - Horizontally flipped image (flip + resize + normalize)\n    \n    These are stacked into shape [2, C, H, W].\n    \n    Args:\n        image_paths: File paths for test images.\n        img_size: Target resolution (must match training).\n    \"\"\"\n    def __init__(self, image_paths: List[Union[str, Path]], img_size: int):\n        self.image_paths = image_paths\n        self.img_size = img_size\n        \n        # Resize + normalize only (no augmentation — test set)\n        self.transform = A.Compose([\n            A.Resize(img_size, img_size),\n            A.Normalize(\n                mean=[0.481, 0.457, 0.408],   # ImageNet statistics\n                std=[0.268, 0.261, 0.275]\n            ),\n        ])\n        \n    def __len__(self) -> int:\n        return len(self.image_paths)\n    \n    def __getitem__(self, idx: int) -> torch.Tensor:\n        path = str(self.image_paths[idx])\n        \n        # Read image (black image on failure)\n        try:\n            image = np.array(Image.open(path).convert(\"RGB\"))\n        except Exception:\n            image = np.zeros((self.img_size, self.img_size, 3), dtype=np.uint8)\n        \n        # Original image → normalize\n        img_orig = self.transform(image=image)[\"image\"]\n        \n        # Horizontal flip (TTA) → normalize\n        img_flip = self.transform(image=cv2.flip(image, 1))[\"image\"]\n        \n        # HWC → CHW conversion and stack [2, C, H, W]\n        return torch.stack([\n            torch.from_numpy(img_orig.transpose(2, 0, 1)).float(),\n            torch.from_numpy(img_flip.transpose(2, 0, 1)).float()\n        ], dim=0)\n\n\ndef extract_features(\n    model: nn.Module, \n    loader: DataLoader, \n    config: Config\n) -> np.ndarray:\n    \"\"\"\n    Extracts L2-normalized embeddings from the model (with TTA averaging).\n    \n    For each batch:\n    1. Flatten [B, 2, C, H, W] → [B×2, C, H, W]\n    2. Produce embeddings via model\n    3. Reshape [B×2, D] → [B, 2, D]\n    4. Average the 2 TTA versions → [B, D]\n    5. L2 normalize\n    \n    Args:\n        model: Inference model (must be in eval mode).\n        loader: Test data loader.\n        config: Configuration (for device info).\n        \n    Returns:\n        L2-normalized embedding matrix [N, D].\n    \"\"\"\n    model.eval()\n    feats = []\n    use_amp = torch.cuda.is_available()  # Mixed precision on CUDA only\n    \n    with torch.no_grad():\n        for imgs in tqdm(loader, desc=\"Extracting Features\"):\n            # imgs shape: [B, T=2, C, H, W] (original + flip)\n            B, T, C, H, W = imgs.shape\n            \n            # Flatten TTA versions into a single batch\n            imgs = imgs.view(-1, C, H, W).to(config.device, non_blocking=True)\n            \n            # Mixed precision inference (if GPU available)\n            if use_amp:\n                with torch.amp.autocast('cuda'):\n                    out = model(imgs)\n            else:\n                out = model(imgs)\n            \n            # [B×T, D] → [B, T, D] → TTA average → [B, D]\n            out = out.view(B, T, -1).mean(dim=1)\n            \n            # L2 normalization → unit vectors\n            out = F.normalize(out, p=2, dim=1)\n            feats.append(out.cpu().numpy())\n    \n    return np.concatenate(feats, axis=0)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-27T18:38:34.876206Z","iopub.execute_input":"2026-03-27T18:38:34.876524Z","iopub.status.idle":"2026-03-27T18:38:34.891217Z","shell.execute_reply.started":"2026-03-27T18:38:34.876501Z","shell.execute_reply":"2026-03-27T18:38:34.890464Z"}},"outputs":[],"execution_count":null},{"id":"0eb66521","cell_type":"markdown","source":"## 6. Main Inference Pipeline\n\nEnd-to-end inference flow bringing all components together:\n\n1. **Data Preparation:** Read test CSV → build unique image list\n2. **Ensemble Inference:** Extract embeddings per checkpoint → weighted sum\n3. **Post-Processing:** L2 Norm → AQE → K-Reciprocal Re-Ranking\n4. **Submission:** Pair similarity scores → gamma calibration → CSV\n","metadata":{}},{"id":"cdfa96a8","cell_type":"code","source":"def main() -> None:\n    \"\"\"\n    Main inference pipeline — ensemble + post-processing + submission generation.\n    \"\"\"\n    config = Config()\n    \n    logger.info(\"=\" * 60)\n    logger.info(\"   Inference Pipeline: Ensemble + AQE + Re-Ranking\")\n    logger.info(\"=\" * 60)\n\n    # =========================================================================\n    # Step 1: Prepare Test Data\n    # =========================================================================\n    df_test = pd.read_csv(config.csv_path)\n    \n    # Get image pair columns (columns 1 and 2)\n    img_col_1, img_col_2 = df_test.columns[1], df_test.columns[2]\n    logger.info(f\"Test pairs: {len(df_test)}\")\n\n    # Collect all unique images and sort\n    all_imgs = sorted(list(set(\n        df_test[img_col_1].tolist() + df_test[img_col_2].tolist()\n    )))\n    logger.info(f\"Unique test images: {len(all_imgs)}\")\n    \n    # Image name → file path mapping\n    img_map = {name: config.test_dir / name for name in all_imgs}\n    \n    # Image name → index mapping (for similarity matrix access)\n    name_to_idx = {name: i for i, name in enumerate(all_imgs)}\n\n    # =========================================================================\n    # Step 2: Create DataLoader (shared across all checkpoints)\n    # =========================================================================\n    dataset = InferenceDataset(list(img_map.values()), config.img_size)\n    loader = DataLoader(\n        dataset, \n        batch_size=config.batch_size, \n        num_workers=config.num_workers, \n        shuffle=False  # Order must be preserved — critical for index mapping\n    )\n\n    # =========================================================================\n    # Step 3: Weighted Feature-Level Ensemble\n    # =========================================================================\n    ensemble_features = None\n    success_count = 0\n\n    for ckpt_name, weight in zip(config.checkpoints, config.weights):\n        weight_path = config.weights_dir / ckpt_name\n        \n        if not weight_path.exists():\n            logger.warning(f\"Skipping (not found): {ckpt_name}\")\n            continue\n\n        logger.info(f\"Processing: {ckpt_name} (weight: {weight:.2f})\")\n        \n        # Create model → load checkpoint → move to GPU\n        model = EVAReIDModelInference(config).to(config.device)\n        model = load_checkpoint(model, weight_path)\n        if model is None:\n            continue\n\n        # Extract embeddings\n        features = extract_features(model, loader, config)\n        \n        # Weighted accumulation (first model → initialize, rest → add)\n        if ensemble_features is None:\n            ensemble_features = features * weight\n        else:\n            ensemble_features += features * weight\n            \n        success_count += 1\n        \n        # Memory cleanup — free space for next model\n        del model\n        gc.collect()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n\n    # =========================================================================\n    # Step 4: Post-Processing Chain\n    # =========================================================================\n    if success_count == 0:\n        logger.error(\"No models loaded successfully! Cannot generate output.\")\n        return\n    \n    logger.info(f\"{success_count} models successfully ensembled.\")\n    \n    # L2 normalization — mandatory after ensemble\n    ensemble_features = ensemble_features / (\n        np.linalg.norm(ensemble_features, axis=1, keepdims=True) + 1e-12\n    )\n\n    # Adaptive Query Expansion\n    if config.use_aqe:\n        logger.info(\n            f\"Applying AQE (threshold: {config.aqe_threshold}, \"\n            f\"max_k: {config.aqe_max_k})...\"\n        )\n        ensemble_features = PostProcessor.adaptive_query_expansion(\n            ensemble_features, \n            threshold=config.aqe_threshold, \n            max_k=config.aqe_max_k\n        )\n\n    # K-Reciprocal Re-Ranking or plain cosine similarity\n    if config.use_re_ranking:\n        logger.info(\n            f\"K-Reciprocal Re-Ranking (k1: {config.re_rank_k1}, \"\n            f\"Jaccard weight: {config.re_rank_weight})...\"\n        )\n        sim_mat = PostProcessor.k_reciprocal_jaccard_similarity(\n            ensemble_features,\n            k1=config.re_rank_k1,\n            jaccard_weight=config.re_rank_weight\n        )\n    else:\n        logger.info(\"Computing cosine similarity matrix...\")\n        sim_mat = np.dot(ensemble_features, ensemble_features.T)\n\n    # =========================================================================\n    # Step 5: Generate Submission File\n    # =========================================================================\n    # Get indices for test pairs\n    idx1_array = np.array([name_to_idx[name] for name in df_test[img_col_1]])\n    idx2_array = np.array([name_to_idx[name] for name in df_test[img_col_2]])\n    \n    # Extract pair scores from similarity matrix and clip to [0, 1]\n    preds = np.clip(sim_mat[idx1_array, idx2_array], 0.0, 1.0)\n    \n    # Gamma calibration: gamma > 1 → makes high scores more prominent\n    if config.calibration_gamma != 1.0:\n        logger.info(f\"Gamma calibration (γ={config.calibration_gamma})...\")\n        preds = np.power(preds, config.calibration_gamma)\n    \n    # Save CSV\n    output_filename = config.output_dir / \"submission.csv\"\n    sub_df = df_test.copy()\n    sub_df[\"similarity\"] = preds\n    sub_df[[\"row_id\", \"similarity\"]].to_csv(output_filename, index=False)\n    \n    logger.info(f\"Submission saved: {output_filename}\")\n    logger.info(f\"Score stats — Min: {preds.min():.4f} | Max: {preds.max():.4f} | Mean: {preds.mean():.4f}\")\n    logger.info(\"Process completed!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-27T18:38:34.892092Z","iopub.execute_input":"2026-03-27T18:38:34.892271Z","iopub.status.idle":"2026-03-27T18:38:34.908504Z","shell.execute_reply.started":"2026-03-27T18:38:34.892252Z","shell.execute_reply":"2026-03-27T18:38:34.907663Z"}},"outputs":[],"execution_count":null},{"id":"c6b1ab86","cell_type":"markdown","source":"## 🚀 Start Inference\n\nRun the cell below to generate the submission file.\n","metadata":{}},{"id":"ad819dc8","cell_type":"code","source":"# =============================================================================\n# Start inference\n# =============================================================================\nif __name__ == \"__main__\":\n    main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-27T18:38:34.909631Z","iopub.execute_input":"2026-03-27T18:38:34.909945Z","iopub.status.idle":"2026-03-27T18:46:02.330412Z","shell.execute_reply.started":"2026-03-27T18:38:34.909909Z","shell.execute_reply":"2026-03-27T18:46:02.329511Z"}},"outputs":[],"execution_count":null}]}