{"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":"!pip install -qU timm\n\nimport os, math, random, warnings, gc\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as T\nimport timm\nwarnings.filterwarnings('ignore')\n\n# ---- Device & Seed ----\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\ndef seed_everything(seed):\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\nseed_everything(42)\n\n# ---- Load Data ----\nINPUT_DIR = '/kaggle/input/jaguar-re-id'\nTRAIN_DIR = os.path.join(INPUT_DIR, 'train/train')\nTEST_DIR = os.path.join(INPUT_DIR, 'test/test')\n\ntrain_df = pd.read_csv(os.path.join(INPUT_DIR, 'train.csv'))\ntest_df = pd.read_csv(os.path.join(INPUT_DIR, 'test.csv'))\nsample_sub = pd.read_csv(os.path.join(INPUT_DIR, 'sample_submission.csv'))\n\n# Label encoding\nlabel_to_idx = {l: i for i, l in enumerate(sorted(train_df['ground_truth'].unique()))}\ntrain_df['label'] = train_df['ground_truth'].map(label_to_idx)\nNUM_CLASSES = len(label_to_idx)\n\nprint(f\"PyTorch {torch.__version__} | Device: {DEVICE}\")\nif torch.cuda.is_available():\n    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\nprint(f\"\\nTraining images: {len(train_df)}\")\nprint(f\"Individuals: {NUM_CLASSES}\")\nprint(f\"Test pairs: {len(test_df):,}\")\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-02-03T20:03:35.517727Z","iopub.execute_input":"2026-02-03T20:03:35.518314Z","iopub.status.idle":"2026-02-03T20:04:00.595914Z","shell.execute_reply.started":"2026-02-03T20:03:35.518288Z","shell.execute_reply":"2026-02-03T20:04:00.595256Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# EXPLORATORY DATA ANALYSIS\n# ============================================================\n\nfig, axes = plt.subplots(4, 5, figsize=(16, 13))\nfig.suptitle('Sample Jaguar Images by Individual', fontsize=16, fontweight='bold')\n\nindividuals = sorted(train_df['ground_truth'].unique())\nnp.random.seed(42)\nselected = list(np.random.choice(individuals, size=min(4, len(individuals)), replace=False))\n\nfor row_idx, indiv in enumerate(selected):\n    indiv_df = train_df[train_df['ground_truth'] == indiv]\n    samples = indiv_df.sample(min(5, len(indiv_df)), random_state=42)\n    for col_idx, (_, img_row) in enumerate(samples.iterrows()):\n        if col_idx >= 5:\n            break\n        try:\n            img = Image.open(os.path.join(TRAIN_DIR, img_row['filename']))\n            axes[row_idx, col_idx].imshow(img)\n        except:\n            pass\n        axes[row_idx, col_idx].axis('off')\n        if col_idx == 0:\n            n_imgs = len(indiv_df)\n            axes[row_idx, col_idx].set_title(f\"{indiv[:20]}\\n({n_imgs} imgs)\", fontsize=9, fontweight='bold')\n    for col_idx in range(len(samples), 5):\n        axes[row_idx, col_idx].axis('off')\n\nplt.tight_layout()\nplt.show()\n\nfig, ax = plt.subplots(figsize=(14, 4))\ncounts = train_df['ground_truth'].value_counts().sort_index()\ncolors = plt.cm.viridis(np.linspace(0.3, 0.9, NUM_CLASSES))\nax.bar(range(NUM_CLASSES), counts.values, color=colors)\nax.set_xlabel('Individual Jaguar')\nax.set_ylabel('Number of Images')\nax.set_title(f'Images per Individual ({len(train_df)} total, {NUM_CLASSES} jaguars)')\nax.set_xticks(range(NUM_CLASSES))\nax.set_xticklabels(counts.index, rotation=45, ha='right', fontsize=7)\nplt.tight_layout()\nplt.show()\n\nprint(f\"Images per individual:\")\nprint(f\"  Min: {counts.min()} | Max: {counts.max()} | Mean: {counts.mean():.1f} | Median: {counts.median():.1f}\")\nprint(f\"\\nTest: {len(set(test_df['query_image']) | set(test_df['gallery_image']))} unique images, {len(test_df):,} pairs\")\nprint(f\"\\nCLOSED-SET: all {NUM_CLASSES} test individuals appear in training.\")\nprint(f\"Strategy: 2x EVA02 (different seeds) + DINOv2, weighted ensemble, QE + re-rank\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 2x EVA02-Large (different seeds) + DINOv2-Base\n# Weighted ensemble + QE + K-Reciprocal Re-ranking\n# ============================================================\n\n# ---- GeM Pooling ----\nclass GeM(nn.Module):\n    def __init__(self, p=3, eps=1e-6):\n        super().__init__()\n        self.p = nn.Parameter(torch.ones(1) * p)\n        self.eps = eps\n    def forward(self, x):\n        return F.avg_pool2d(x.clamp(min=self.eps).pow(self.p),\n                           (x.size(-2), x.size(-1))).pow(1.0 / self.p)\n\n# ---- ArcFace ----\nclass ArcFaceLayer(nn.Module):\n    def __init__(self, in_features, out_features, s=30.0, m=0.5):\n        super().__init__()\n        self.s, self.m = s, m\n        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))\n        nn.init.xavier_uniform_(self.weight)\n    def forward(self, input, label=None):\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n        if label is None:\n            return cosine\n        phi = cosine - self.m\n        one_hot = torch.zeros_like(cosine)\n        one_hot.scatter_(1, label.view(-1, 1), 1)\n        return ((one_hot * phi) + ((1.0 - one_hot) * cosine)) * self.s\n\n# ---- Model (EVA02 + ViT compatible) ----\nclass JaguarModel(nn.Module):\n    def __init__(self, model_name, num_classes, s=30.0, m=0.5):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=True,\n                                          num_classes=0, dynamic_img_size=True)\n        self.feat_dim = self.backbone.num_features\n        self.gem = GeM()\n        self.bn = nn.BatchNorm1d(self.feat_dim)\n        self.head = ArcFaceLayer(self.feat_dim, num_classes, s=s, m=m)\n        n_params = sum(p.numel() for p in self.parameters())\n        print(f\"  Dim: {self.feat_dim} | Params: {n_params:,}\")\n\n    def forward(self, x, label=None):\n        features = self.backbone.forward_features(x)\n\n        if features.dim() == 3:\n            B, N, C = features.shape\n            H = W = int(math.sqrt(N))\n            if H * W != N:\n                features = features[:, -H * W:, :]\n            features = features.permute(0, 2, 1).reshape(B, C, H, W)\n\n        emb = self.gem(features).flatten(1)\n        emb = self.bn(emb)\n        if label is not None:\n            return self.head(emb, label)\n        return emb\n\n# ---- Dataset ----\nclass JaguarDataset(Dataset):\n    def __init__(self, df, img_dir, transform=None, is_test=False):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = img_dir\n        self.transform = transform\n        self.is_test = is_test\n    def __len__(self):\n        return len(self.df)\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        try:\n            img = Image.open(os.path.join(self.img_dir, row['filename'])).convert('RGB')\n        except Exception:\n            img = Image.new('RGB', (448, 448))\n        if self.transform:\n            img = self.transform(img)\n        if self.is_test:\n            return img, row['filename']\n        return img, torch.tensor(row['label'], dtype=torch.long)\n\n# ---- Transforms ----\ndef get_train_transform(img_size, strength='normal'):\n    if strength == 'strong':\n        return T.Compose([\n            T.Resize((img_size, img_size)),\n            T.RandomHorizontalFlip(),\n            T.RandomAffine(degrees=20, translate=(0.15, 0.15), scale=(0.85, 1.15)),\n            T.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2),\n            T.ToTensor(),\n            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n            T.RandomErasing(p=0.35),\n        ])\n    return T.Compose([\n        T.Resize((img_size, img_size)),\n        T.RandomHorizontalFlip(),\n        T.RandomAffine(degrees=15, translate=(0.1, 0.1), scale=(0.9, 1.1)),\n        T.ColorJitter(brightness=0.2, contrast=0.2),\n        T.ToTensor(),\n        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n        T.RandomErasing(p=0.25),\n    ])\n\ndef get_test_transform(img_size):\n    return T.Compose([\n        T.Resize((img_size, img_size)),\n        T.ToTensor(),\n        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n    ])\n\n# ============================================================\n# TRAINING FUNCTION\n# ============================================================\ndef train_and_extract(model_name, img_size, batch_size, grad_accum, epochs, lr, wd,\n                      seed=42, aug_strength='normal'):\n    seed_everything(seed)\n\n    print(f\"\\n{'='*60}\")\n    print(f\"TRAINING: {model_name.split('.')[-1]} (seed={seed}, aug={aug_strength})\")\n    print(f\"  {img_size}px | bs={batch_size}x{grad_accum} | {epochs} epochs | lr={lr}\")\n    print(f\"{'='*60}\")\n\n    model = JaguarModel(model_name, NUM_CLASSES).to(DEVICE)\n\n    train_loader = DataLoader(\n        JaguarDataset(train_df.copy(), TRAIN_DIR, get_train_transform(img_size, aug_strength)),\n        batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True\n    )\n\n    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n    scaler = torch.amp.GradScaler('cuda')\n    criterion = nn.CrossEntropyLoss()\n\n    for epoch in range(epochs):\n        model.train()\n        total_loss = 0\n        optimizer.zero_grad()\n        for i, (imgs, labels) in enumerate(tqdm(train_loader, leave=False, desc=f\"Epoch {epoch+1}\")):\n            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n            with torch.amp.autocast('cuda'):\n                loss = criterion(model(imgs, labels), labels) / grad_accum\n            scaler.scale(loss).backward()\n            if (i + 1) % grad_accum == 0:\n                scaler.step(optimizer)\n                scaler.update()\n                optimizer.zero_grad()\n            total_loss += loss.item() * grad_accum\n        # Flush remaining gradients\n        scaler.step(optimizer)\n        scaler.update()\n        optimizer.zero_grad()\n        scheduler.step()\n        print(f\"  Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(train_loader):.4f} | LR: {scheduler.get_last_lr()[0]:.2e}\")\n\n    # Extract embeddings with TTA (original + hflip)\n    print(\"  Extracting embeddings...\")\n    test_loader = DataLoader(\n        JaguarDataset(pd.DataFrame({'filename': unique_test}), TEST_DIR,\n                     get_test_transform(img_size), is_test=True),\n        batch_size=batch_size, shuffle=False, num_workers=2\n    )\n\n    model.eval()\n    feats, names = [], []\n    with torch.no_grad():\n        for imgs, fnames in tqdm(test_loader, leave=False, desc=\"Inference\"):\n            imgs = imgs.to(DEVICE)\n            with torch.amp.autocast('cuda'):\n                f1 = model(imgs)\n                f2 = model(torch.flip(imgs, [3]))\n            f_avg = F.normalize((f1 + f2) / 2, dim=1)\n            feats.append(f_avg.float().cpu())\n            names.extend(fnames)\n\n    embeddings = torch.cat(feats, 0).numpy()\n\n    del model, optimizer, scaler, scheduler, train_loader, test_loader\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    print(f\"  Done! Embeddings: {embeddings.shape}\")\n    return embeddings, names\n\n# ============================================================\n# POST-PROCESSING\n# ============================================================\ndef query_expansion(emb, top_k=3):\n    print(\"  Applying Query Expansion (top_k=3)...\")\n    sims = emb @ emb.T\n    indices = np.argsort(-sims, axis=1)[:, :top_k]\n    new_emb = np.zeros_like(emb)\n    for i in range(len(emb)):\n        new_emb[i] = np.mean(emb[indices[i]], axis=0)\n    return new_emb / np.linalg.norm(new_emb, axis=1, keepdims=True)\n\ndef k_reciprocal_rerank(prob, k1=20, k2=6, lambda_value=0.3):\n    print(\"  Applying K-Reciprocal Re-ranking...\")\n    q_g_dist = 1 - prob\n    original_dist = q_g_dist.copy()\n    initial_rank = np.argsort(original_dist, axis=1)\n\n    nn_k1 = []\n    for i in range(prob.shape[0]):\n        forward_k1 = initial_rank[i, :k1 + 1]\n        backward_k1 = initial_rank[forward_k1, :k1 + 1]\n        fi = np.where(backward_k1 == i)[0]\n        nn_k1.append(forward_k1[fi])\n\n    jaccard_dist = np.zeros_like(original_dist)\n    for i in range(prob.shape[0]):\n        ind_non_zero = np.where(original_dist[i, :] < 0.6)[0]\n        ind_images = [inv for inv in ind_non_zero\n                      if len(np.intersect1d(nn_k1[i], nn_k1[inv])) > 0]\n        for j in ind_images:\n            intersection = len(np.intersect1d(nn_k1[i], nn_k1[j]))\n            union = len(np.union1d(nn_k1[i], nn_k1[j]))\n            jaccard_dist[i, j] = 1 - intersection / union\n\n    return 1 - (jaccard_dist * lambda_value + original_dist * (1 - lambda_value))\n\n# ============================================================\n# MAIN: TRAIN 3 MODELS, WEIGHTED ENSEMBLE, SUBMIT\n# ============================================================\nunique_test = sorted(set(test_df['query_image']) | set(test_df['gallery_image']))\nprint(f\"\\nUnique test images: {len(unique_test)}\")\nprint(f\"\\n{'#'*60}\")\nprint(f\"# STRATEGY: Beat 0.937\")\nprint(f\"# 1. EVA02-Large seed=42  (normal aug)  -> weight 0.40\")\nprint(f\"# 2. EVA02-Large seed=123 (strong aug)  -> weight 0.40\")\nprint(f\"# 3. DINOv2-Base seed=42  (normal aug)  -> weight 0.20\")\nprint(f\"# Weighted ensemble -> QE -> Re-ranking\")\nprint(f\"{'#'*60}\")\n\nMODELS = [\n    # (name, img_size, bs, ga, epochs, lr, wd, seed, aug, weight)\n    ('eva02_large_patch14_448.mim_m38m_ft_in22k_in1k', 448, 4, 4, 10, 2e-5, 1e-3, 42,  'normal', 0.40),\n    ('eva02_large_patch14_448.mim_m38m_ft_in22k_in1k', 448, 4, 4, 10, 2e-5, 1e-3, 123, 'strong', 0.40),\n    ('vit_base_patch14_dinov2.lvd142m',                 518, 8, 2, 10, 1e-5, 1e-4, 42,  'normal', 0.20),\n]\n\nall_sim_matrices = []\nall_weights = []\nimg_map = None\n\nfor model_name, img_size, bs, ga, epochs, lr, wd, seed, aug, weight in MODELS:\n    try:\n        emb, names = train_and_extract(model_name, img_size, bs, ga, epochs, lr, wd,\n                                       seed=seed, aug_strength=aug)\n        if img_map is None:\n            img_map = {n: i for i, n in enumerate(names)}\n\n        emb = query_expansion(emb)\n        sim = emb @ emb.T\n        all_sim_matrices.append(sim)\n        all_weights.append(weight)\n        print(f\"  Model complete! (weight={weight}) \\u2705\")\n\n    except Exception as e:\n        print(f\"\\n  FAILED: {e}\")\n        import traceback\n        traceback.print_exc()\n        print(f\"  Skipping...\\n\")\n        gc.collect()\n        torch.cuda.empty_cache()\n        continue\n\n# ---- Weighted Ensemble ----\nprint(f\"\\n{'='*60}\")\nprint(f\"WEIGHTED ENSEMBLE: {len(all_sim_matrices)} models\")\nprint(f\"{'='*60}\")\n\nif len(all_sim_matrices) == 0:\n    raise RuntimeError(\"No models trained successfully!\")\n\n# Normalize weights to sum to 1.0\ntotal_w = sum(all_weights)\nnorm_weights = [w / total_w for w in all_weights]\nprint(f\"  Weights: {norm_weights}\")\n\ncombined_sim = sum(w * s for w, s in zip(norm_weights, all_sim_matrices))\ncombined_sim = k_reciprocal_rerank(combined_sim)\n\n# ---- Generate Submission ----\nprint(\"\\nGenerating submission...\")\npreds = []\nfor _, row in tqdm(test_df.iterrows(), total=len(test_df), desc=\"Mapping\"):\n    s = combined_sim[img_map[row['query_image']], img_map[row['gallery_image']]]\n    preds.append(max(0.0, min(1.0, float(s))))\n\nsub = pd.DataFrame({'row_id': test_df['row_id'], 'similarity': preds})\nsub.to_csv('submission.csv', index=False)\n\nprint(f\"\\n{'='*60}\")\nprint(f\"DONE!\")\nprint(f\"{'='*60}\")\nprint(f\"Models ensembled: {len(all_sim_matrices)}\")\nprint(f\"Weights: {norm_weights}\")\nprint(f\"Mean similarity: {np.mean(preds):.4f}\")\nprint(f\"Range: [{min(preds):.4f}, {max(preds):.4f}]\")\nprint(f\"\\nSubmission saved to submission.csv\")\nprint(sub.head(10))\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}