{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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":[{"sourceId":126777,"databundleVersionId":15314950,"sourceType":"competition"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-23T20:02:04.422993Z","iopub.execute_input":"2026-01-23T20:02:04.423549Z","iopub.status.idle":"2026-01-23T20:02:08.914871Z","shell.execute_reply.started":"2026-01-23T20:02:04.423525Z","shell.execute_reply":"2026-01-23T20:02:08.914105Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🐆 Jaguar Re-ID (Swin + Metric Learning) — End-to-End Notebook\n\n**Objective:** learn an embedding space where images of the same jaguar are close and different jaguars are far, then predict **pairwise similarity scores** for all test image pairs (137,270 rows) required by the competition.\n\nThis notebook trains a **Swin** backbone with a **margin-softmax (ArcFace-style) head**, extracts embeddings for all test images once, computes cosine similarities, and optionally applies **TTA (flip)** and **k-reciprocal re-ranking** to improve retrieval rankings.\n","metadata":{}},{"cell_type":"markdown","source":"## Method overview (high-level)\n\n**Training**\n1. Load labeled training cutouts.\n2. Train an embedding model:\n   - Backbone: `timm` Swin (fine-grained features)\n   - Projection head: BN → Linear → BN (embedding)\n   - Classification head: margin-softmax (ArcFace-style) over identities\n\n**Inference**\n1. Extract **L2-normalized embeddings** for all unique test images (371).\n2. Compute cosine similarity matrix `S = E @ Eᵀ`.\n3. Optional:\n   - **TTA**: average embeddings from original + horizontal flip\n   - **Re-ranking**: k-reciprocal neighborhood consistency (often improves retrieval)\n4. Map similarities to `[0,1]`, then fill `test.csv` pairs in exact row order and save `submission.csv`.\n","metadata":{}},{"cell_type":"code","source":"# ============================================================\n# IMPORTS\n# ============================================================\n\nimport os\nimport gc\nimport math\nimport time\nimport random\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\n\nfrom PIL import Image\n\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\n\nimport torchvision.transforms as transforms\n\nfrom sklearn.preprocessing import LabelEncoder\n\nimport timm\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T23:01:36.776460Z","iopub.execute_input":"2026-01-23T23:01:36.776743Z","iopub.status.idle":"2026-01-23T23:01:36.781792Z","shell.execute_reply.started":"2026-01-23T23:01:36.776719Z","shell.execute_reply":"2026-01-23T23:01:36.781122Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Configuration\n\nKey knobs:\n- `mode`: `\"fast\"` for debugging vs `\"full\"` for best score\n- `backbone`: Swin variants from `timm` (better for fine-grained patterns)\n- `img_size`: 224 is faster; 384 can help if GPU allows\n- `arcface_s`, `arcface_m`: scale + margin for margin-softmax head\n- `use_tta`: flip augmentation at inference (more stable embeddings)\n- `use_reranking`: k-reciprocal re-ranking (improves ranking quality; slower)\n","metadata":{}},{"cell_type":"code","source":"\n# ============================================================================\n# CONFIGURATION\n# ============================================================================\nclass Config:\n    # Mode: \"fast\" or \"full\"\n    mode = \"full\"\n    \n    # Paths\n    data_root = \"/kaggle/input/jaguar-re-id\"\n    \n    # Model - Swin is better for fine-grained patterns\n    backbone = \"swin_base_patch4_window7_224\"\n    img_size = 224\n    embedding_dim = 512\n    \n    # Training\n    batch_size = 32\n    num_epochs = 20\n    lr = 1e-4\n    weight_decay = 0.01\n    \n    # ArcFace (simple version that works)\n    arcface_s = 64.0\n    arcface_m = 0.5\n    \n    # Inference improvements\n    use_tta = True           # Test-time augmentation\n    use_reranking = True     # K-reciprocal re-ranking\n    rerank_k1 = 20\n    rerank_k2 = 6\n    rerank_lambda = 0.3\n    \n    # System\n    seed = 42\n    num_workers = 2\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n# Fast mode adjustments\nif Config.mode == \"fast\":\n    Config.num_epochs = 10\n    Config.use_reranking = False\n    Config.use_tta = False\n\ndef seed_everything(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nseed_everything(Config.seed)\n\nprint(\"=\"*70)\nprint(f\"JAGUAR RE-ID NOTEBOOK | mode={Config.mode.upper()}\")\nprint(\"=\"*70)\nprint(f\"data_root     : {Config.data_root}\")\nprint(f\"device        : {Config.device}\")\nprint(f\"backbone      : {Config.backbone}\")\nprint(f\"img_size      : {Config.img_size}\")\nprint(f\"embedding_dim : {Config.embedding_dim}\")\nprint(f\"epochs        : {Config.num_epochs}\")\nprint(f\"batch_size    : {Config.batch_size}\")\nprint(f\"lr / wd       : {Config.lr} / {Config.weight_decay}\")\nprint(f\"TTA           : {Config.use_tta}\")\nprint(f\"re-ranking    : {Config.use_reranking} (k1={Config.rerank_k1}, k2={Config.rerank_k2}, λ={Config.rerank_lambda})\")\nprint(\"=\"*70)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T23:01:39.166363Z","iopub.execute_input":"2026-01-23T23:01:39.167037Z","iopub.status.idle":"2026-01-23T23:01:39.177971Z","shell.execute_reply.started":"2026-01-23T23:01:39.167003Z","shell.execute_reply":"2026-01-23T23:01:39.177419Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# DATA LOADING\n# ============================================================================\nDATA_ROOT = Path(Config.data_root)\nTRAIN_DIR = DATA_ROOT / \"train\" / \"train\"\nTEST_DIR = DATA_ROOT / \"test\" / \"test\"\n\ntrain_df = pd.read_csv(DATA_ROOT / \"train.csv\")\ntest_df = pd.read_csv(DATA_ROOT / \"test.csv\")\n\n# Find column names dynamically\nid_col = [c for c in train_df.columns if c not in ['image', 'filename']][0]\nimg_col = 'image' if 'image' in train_df.columns else 'filename'\n\nprint(f\"\\nID column: '{id_col}'\")\nprint(f\"Image column: '{img_col}'\")\nprint(f\"Training images: {len(train_df)}\")\nprint(f\"Test pairs: {len(test_df)}\")\n\n# Create label mapping\nunique_ids = sorted(train_df[id_col].unique())\nlabel_map = {jag_id: idx for idx, jag_id in enumerate(unique_ids)}\ntrain_df['label'] = train_df[id_col].map(label_map)\nnum_classes = len(unique_ids)\n\nprint(f\"Unique jaguars: {num_classes}\")\n\n# Class distribution\nprint(\"\\nClass distribution:\")\nclass_counts = train_df[id_col].value_counts()\nprint(f\"  Max: {class_counts.max()} ({class_counts.idxmax()})\")\nprint(f\"  Min: {class_counts.min()} ({class_counts.idxmin()})\")\nprint(f\"  Mean: {class_counts.mean():.1f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T23:01:40.543904Z","iopub.execute_input":"2026-01-23T23:01:40.544657Z","iopub.status.idle":"2026-01-23T23:01:40.624544Z","shell.execute_reply.started":"2026-01-23T23:01:40.544629Z","shell.execute_reply":"2026-01-23T23:01:40.623944Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"\\n[Sanity checks]\")\nprint(\"train_df:\", train_df.shape)\nprint(\"test_df :\", test_df.shape)\nprint(\"num_classes:\", num_classes)\nprint(\"train dir exists:\", TRAIN_DIR.exists(), \"|\", TRAIN_DIR)\nprint(\"test dir exists :\", TEST_DIR.exists(),  \"|\", TEST_DIR)\n\n# Check 3 sample files exist\nfor fn in train_df[img_col].head(3).tolist():\n    p = TRAIN_DIR / fn\n    print(\"exists:\", p.exists(), \"-\", p)\n\n# Check unique test images count\nunique_test_imgs = sorted(set(test_df[\"query_image\"]) | set(test_df[\"gallery_image\"]))\nprint(\"unique test imgs:\", len(unique_test_imgs))\nassert len(unique_test_imgs) == 371, \"Expected 371 unique test images\"\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T23:01:41.667511Z","iopub.execute_input":"2026-01-23T23:01:41.668013Z","iopub.status.idle":"2026-01-23T23:01:41.696079Z","shell.execute_reply.started":"2026-01-23T23:01:41.667988Z","shell.execute_reply":"2026-01-23T23:01:41.695566Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# DATASET\n# ============================================================================\nclass JaguarDataset(Dataset):\n    \"\"\"\n    Loads jaguar cutout images and returns:\n      - train mode: (image_tensor, label)\n      - test mode : (image_tensor, filename)\n    Notes:\n    - Expects img_dir / filename to exist.\n    - Uses PIL -> RGB -> torchvision transforms.\n    \"\"\"\n    def __init__(self, df, img_dir, img_col, transform=None, is_train=True):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = Path(img_dir)\n        self.img_col = img_col\n        self.transform = transform\n        self.is_train = is_train\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = self.img_dir / row[self.img_col]\n        img = Image.open(img_path).convert('RGB')\n        \n        if self.transform:\n            img = self.transform(img)\n        \n        if self.is_train:\n            return img, torch.tensor(row['label'], dtype=torch.long)\n        else:\n            return img, row[self.img_col]\n\n# Transforms\ntrain_transform = transforms.Compose([\n    transforms.Resize((Config.img_size, Config.img_size)),\n    transforms.RandomHorizontalFlip(p=0.5),\n    transforms.RandomRotation(15),\n    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n])\n\ntest_transform = transforms.Compose([\n    transforms.Resize((Config.img_size, Config.img_size)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T23:03:08.565768Z","iopub.execute_input":"2026-01-23T23:03:08.566447Z","iopub.status.idle":"2026-01-23T23:03:08.573816Z","shell.execute_reply.started":"2026-01-23T23:03:08.566419Z","shell.execute_reply":"2026-01-23T23:03:08.573105Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# ARCFACE (Simple version that works)\n# ============================================================================\nclass ArcFaceLayer(nn.Module):\n    def __init__(self, in_features, out_features, s=64.0, m=0.5):\n        super().__init__()\n        self.s = s\n        self.m = m\n        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))\n        nn.init.xavier_uniform_(self.weight)\n\n    def forward(self, embeddings, labels=None):\n        # Normalize\n        embeddings = F.normalize(embeddings, dim=1)\n        weight = F.normalize(self.weight, dim=1)\n        \n        # Cosine similarity\n        cosine = F.linear(embeddings, weight)\n        \n        if labels is None:\n            return cosine\n        \n        # Subtract margin from target class\n        one_hot = torch.zeros_like(cosine)\n        one_hot.scatter_(1, labels.view(-1, 1), 1)\n        output = cosine - one_hot * self.m\n        output *= self.s\n        \n        return output\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T23:03:15.724725Z","iopub.execute_input":"2026-01-23T23:03:15.725454Z","iopub.status.idle":"2026-01-23T23:03:15.731363Z","shell.execute_reply.started":"2026-01-23T23:03:15.725425Z","shell.execute_reply":"2026-01-23T23:03:15.730490Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# MODEL\n# ============================================================================\nclass JaguarReIDModel(nn.Module):\n    def __init__(self, backbone_name, embedding_dim, num_classes, pretrained=True):\n        super().__init__()\n        self.backbone = timm.create_model(backbone_name, pretrained=pretrained, num_classes=0)\n        \n        # Get backbone output dimension\n        with torch.no_grad():\n            dummy = torch.randn(1, 3, Config.img_size, Config.img_size)\n            backbone_dim = self.backbone(dummy).shape[1]\n        \n        # Projection head\n        self.neck = nn.Sequential(\n            nn.BatchNorm1d(backbone_dim),\n            nn.Dropout(0.2),\n            nn.Linear(backbone_dim, embedding_dim),\n            nn.BatchNorm1d(embedding_dim),\n        )\n        \n        # ArcFace head\n        self.head = ArcFaceLayer(embedding_dim, num_classes, s=Config.arcface_s, m=Config.arcface_m)\n        \n    def forward(self, x, labels=None):\n        features = self.backbone(x)\n        embeddings = self.neck(features)\n        \n        if labels is not None:\n            logits = self.head(embeddings, labels)\n            return logits, embeddings\n        \n        return F.normalize(embeddings, dim=1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T23:03:17.824275Z","iopub.execute_input":"2026-01-23T23:03:17.824604Z","iopub.status.idle":"2026-01-23T23:03:17.830885Z","shell.execute_reply.started":"2026-01-23T23:03:17.824581Z","shell.execute_reply":"2026-01-23T23:03:17.830201Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# TRAINING\n# ============================================================================\ndef train_epoch(model, loader, criterion, optimizer, scaler, device):\n    \"\"\"\n    One training epoch with AMP.\n    Returns average loss over the epoch.\n    \"\"\"\n    model.train()\n    total_loss = 0\n    \n    pbar = tqdm(loader, desc=\"Training\")\n    for images, labels in pbar:\n        images, labels = images.to(device), labels.to(device)\n        \n        optimizer.zero_grad()\n        \n        with autocast():\n            logits, _ = model(images, labels)\n            loss = criterion(logits, labels)\n        \n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        \n        total_loss += loss.item()\n        pbar.set_postfix({'loss': f'{loss.item():.4f}'})\n    \n    return total_loss / len(loader)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T23:03:19.360878Z","iopub.execute_input":"2026-01-23T23:03:19.361619Z","iopub.status.idle":"2026-01-23T23:03:19.367136Z","shell.execute_reply.started":"2026-01-23T23:03:19.361590Z","shell.execute_reply":"2026-01-23T23:03:19.366419Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# INFERENCE WITH TTA\n# ============================================================================\n@torch.no_grad()\ndef extract_embeddings(model, loader, device, use_tta=False):\n    \"\"\"\n    Extract L2-normalized embeddings for a set of images.\n    If use_tta=True: average embedding(original) and embedding(horizontal_flip).\n    Returns: dict {filename: embedding_vector (np.float32)}\n    \"\"\"\n    model.eval()\n    all_embeddings = []\n    all_names = []\n    \n    for images, names in tqdm(loader, desc=\"Extracting embeddings\"):\n        images = images.to(device)\n        \n        if use_tta:\n            # Original + horizontal flip\n            emb1 = model(images)\n            emb2 = model(torch.flip(images, dims=[3]))\n            embeddings = F.normalize((emb1 + emb2) / 2, dim=1)\n        else:\n            embeddings = model(images)\n        \n        all_embeddings.append(embeddings.cpu().numpy())\n        all_names.extend(names)\n    \n    return dict(zip(all_names, np.vstack(all_embeddings)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T23:03:22.595609Z","iopub.execute_input":"2026-01-23T23:03:22.596315Z","iopub.status.idle":"2026-01-23T23:03:22.601616Z","shell.execute_reply.started":"2026-01-23T23:03:22.596287Z","shell.execute_reply":"2026-01-23T23:03:22.601017Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Optional: k-reciprocal re-ranking\n\nRe-ranking adjusts pairwise similarities using neighborhood agreement:\n- If A is among B’s nearest neighbors AND B is among A’s nearest neighbors, they reinforce each other.\n- This often improves retrieval metrics (like mAP) by fixing ambiguous rankings.\n\nThis step is optional because it adds computation, but with only 371 test images it is feasible.\n","metadata":{}},{"cell_type":"code","source":"# ============================================================================\n# K-RECIPROCAL RE-RANKING\n# ============================================================================\ndef k_reciprocal_rerank(sim_matrix, k1=20, k2=6, lambda_val=0.3):\n    \"\"\"K-reciprocal re-ranking - typically adds +3-5% mAP in Re-ID.\"\"\"\n    n = sim_matrix.shape[0]\n    dist_matrix = 1 - sim_matrix\n    initial_rank = np.argsort(dist_matrix, axis=1)\n    \n    # Build k-reciprocal neighbor sets\n    k_reciprocal_sets = []\n    for i in range(n):\n        k_neighbors = set(initial_rank[i, :k1+1].tolist())\n        k_recip = set()\n        for j in k_neighbors:\n            if i in set(initial_rank[j, :k1+1].tolist()):\n                k_recip.add(j)\n        k_reciprocal_sets.append(k_recip)\n    \n    # Compute Jaccard distance\n    jaccard = np.zeros((n, n), dtype=np.float32)\n    for i in range(n):\n        for j in range(n):\n            if i != j:\n                inter = len(k_reciprocal_sets[i] & k_reciprocal_sets[j])\n                union = len(k_reciprocal_sets[i] | k_reciprocal_sets[j])\n                jaccard[i, j] = 1 - inter / union if union > 0 else 1\n    \n    # Combine distances\n    final_dist = lambda_val * jaccard + (1 - lambda_val) * dist_matrix\n    return 1 - final_dist","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T22:53:28.736832Z","iopub.execute_input":"2026-01-23T22:53:28.737451Z","iopub.status.idle":"2026-01-23T22:53:28.743498Z","shell.execute_reply.started":"2026-01-23T22:53:28.737426Z","shell.execute_reply":"2026-01-23T22:53:28.742918Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# MAIN PIPELINE\n# ============================================================================\nprint(\"\\n\" + \"=\"*60)\nprint(\"🔄 TRAINING PHASE (using ALL data)\")\nprint(\"=\"*60)\n\n# Create dataset - USE ALL DATA (no validation split)\ntrain_dataset = JaguarDataset(train_df, TRAIN_DIR, img_col, train_transform, is_train=True)\ntrain_loader = DataLoader(\n    train_dataset, \n    batch_size=Config.batch_size, \n    shuffle=True, \n    num_workers=Config.num_workers,\n    pin_memory=True\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T22:53:31.229130Z","iopub.execute_input":"2026-01-23T22:53:31.229447Z","iopub.status.idle":"2026-01-23T22:53:31.235377Z","shell.execute_reply.started":"2026-01-23T22:53:31.229423Z","shell.execute_reply":"2026-01-23T22:53:31.234632Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Initialize model\nmodel = JaguarReIDModel(Config.backbone, Config.embedding_dim, num_classes).to(Config.device)\ncriterion = nn.CrossEntropyLoss()\noptimizer = torch.optim.AdamW(model.parameters(), lr=Config.lr, weight_decay=Config.weight_decay)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=Config.num_epochs)\nscaler = GradScaler()\n\nprint(f\"Model parameters: {sum(p.numel() for p in model.parameters()):,}\")\n\n# Training loop\nbest_loss = float('inf')\nfor epoch in range(Config.num_epochs):\n    print(f\"\\nEpoch {epoch+1}/{Config.num_epochs}\")\n    \n    loss = train_epoch(model, train_loader, criterion, optimizer, scaler, Config.device)\n    scheduler.step()\n    \n    print(f\"Loss: {loss:.4f} | LR: {scheduler.get_last_lr()[0]:.6f}\")\n    \n    if loss < best_loss:\n        best_loss = loss\n        torch.save(model.state_dict(), 'best_model.pth')\n        print(\"💾 Saved best model!\")\n\n# Load best model\nprint(\"\\n📥 Loading best model...\")\nmodel.load_state_dict(torch.load('best_model.pth'))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Inference (submission generation)\n\nWe:\n1. List all unique test images (371).\n2. Extract embeddings once (optionally with TTA).\n3. Build a similarity matrix via cosine similarity (`E @ Eᵀ`).\n4. (Optional) apply re-ranking.\n5. Convert to `[0,1]` and fill `test.csv` pairs in **exact row order**.\n6. Save `submission.csv`.\n","metadata":{}},{"cell_type":"code","source":"# ============================================================================\n# INFERENCE PHASE\n# ============================================================================\nprint(\"\\n\" + \"=\"*60)\nprint(\"🔍 INFERENCE PHASE\")\nprint(\"=\"*60)\n\n# Get unique test images\nunique_test_imgs = sorted(set(test_df['query_image']) | set(test_df['gallery_image']))\nprint(f\"Unique test images: {len(unique_test_imgs)}\")\n\ntest_imgs_df = pd.DataFrame({img_col: unique_test_imgs})\ntest_dataset = JaguarDataset(test_imgs_df, TEST_DIR, img_col, test_transform, is_train=False)\ntest_loader = DataLoader(\n    test_dataset, \n    batch_size=Config.batch_size, \n    shuffle=False, \n    num_workers=Config.num_workers\n)\n\n# Extract embeddings (with TTA if enabled)\nembeddings_dict = extract_embeddings(model, test_loader, Config.device, use_tta=Config.use_tta)\n\n# Build embedding matrix\nemb_matrix = np.array([embeddings_dict[img] for img in unique_test_imgs])\nprint(f\"Embedding matrix shape: {emb_matrix.shape}\")\n\n# Compute similarity matrix\nprint(\"\\n🔢 Computing similarities...\")\nsim_matrix = emb_matrix @ emb_matrix.T\n\n# Apply re-ranking if enabled\nif Config.use_reranking:\n    print(\"🔄 Applying k-reciprocal re-ranking...\")\n    sim_matrix = k_reciprocal_rerank(\n        sim_matrix, \n        k1=Config.rerank_k1, \n        k2=Config.rerank_k2, \n        lambda_val=Config.rerank_lambda\n    )\n\n# Normalize to [0, 1]\nsim_matrix = (sim_matrix + 1) / 2\nsim_matrix = np.clip(sim_matrix, 0, 1)\n\n# Build filename to index mapping\nimg_to_idx = {img: i for i, img in enumerate(unique_test_imgs)}\n\n# Generate similarities for all pairs\nprint(\" Generating submission...\")\nsimilarities = []\nfor _, row in tqdm(test_df.iterrows(), total=len(test_df)):\n    q_idx = img_to_idx[row['query_image']]\n    g_idx = img_to_idx[row['gallery_image']]\n    similarities.append(sim_matrix[q_idx, g_idx])\n\nsimilarities = np.array(similarities)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# SAVE SUBMISSION\n# ============================================================================\nprint(\"\\n\" + \"=\"*60)\nprint(\"💾 SAVING SUBMISSION\")\nprint(\"=\"*60)\n\nsubmission = pd.DataFrame({\n    'row_id': test_df['row_id'],\n    'similarity': similarities\n})\n\nsubmission.to_csv('submission.csv', index=False)\n\nprint(f\"✅ Submission saved: {len(submission)} predictions\")\nprint(f\"Similarity range: [{similarities.min():.4f}, {similarities.max():.4f}]\")\nprint(f\"Mean: {similarities.mean():.4f}, Std: {similarities.std():.4f}\")\n\n# Sanity checks\nassert len(submission) == 137270, f\"Wrong length: {len(submission)}\"\nassert submission['similarity'].notna().all(), \"Contains NaN!\"\nassert (submission['similarity'] >= 0).all() and (submission['similarity'] <= 1).all(), \"Out of range!\"\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"🎯 COMPLETE!\")\nprint(\"=\"*60)\nprint(f\"Model: {Config.backbone}\")\nprint(f\"Epochs: {Config.num_epochs}\")\nprint(f\"TTA: {Config.use_tta}\")\nprint(f\"Re-ranking: {Config.use_reranking}\")\nprint(\"\\n🚀 Ready to submit!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T23:12:43.068042Z","iopub.execute_input":"2026-01-23T23:12:43.068688Z","iopub.status.idle":"2026-01-23T23:12:43.260501Z","shell.execute_reply.started":"2026-01-23T23:12:43.068660Z","shell.execute_reply":"2026-01-23T23:12:43.259636Z"}},"outputs":[],"execution_count":null}]}