{"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":"none","dataSources":[{"sourceType":"competition","sourceId":126777,"databundleVersionId":15314950}],"dockerImageVersionId":31286,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"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 if there wsndjfn jdnnfdudn ncn  jcjiiid dbgcjudn","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:26:39.834646Z","iopub.execute_input":"2026-04-02T03:26:39.83515Z","iopub.status.idle":"2026-04-02T03:26:45.133515Z","shell.execute_reply.started":"2026-04-02T03:26:39.835096Z","shell.execute_reply":"2026-04-02T03:26:45.132905Z"}},"outputs":[],"execution_count":null},{"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\nfrom sklearn.model_selection import StratifiedKFold\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\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):,}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:26:45.134977Z","iopub.execute_input":"2026-04-02T03:26:45.135274Z","iopub.status.idle":"2026-04-02T03:27:04.578175Z","shell.execute_reply.started":"2026-04-02T03:26:45.135255Z","shell.execute_reply":"2026-04-02T03:27:04.577233Z"}},"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'])).convert('RGB')\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            axes[row_idx, col_idx].set_title(f\"{indiv[:20]}\\n({len(indiv_df)} imgs)\", fontsize=9, fontweight='bold')\n    for col_idx in range(len(samples), 5):\n        axes[row_idx, col_idx].axis('off')\nplt.tight_layout()\nplt.show()\n\nfig, ax = plt.subplots(figsize=(14, 4))\ncounts = train_df['ground_truth'].value_counts().sort_index()\nax.bar(range(NUM_CLASSES), counts.values, color=plt.cm.viridis(np.linspace(0.3, 0.9, NUM_CLASSES)))\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: Min={counts.min()} Max={counts.max()} Mean={counts.mean():.1f}\")\nprint(f\"Test: {len(set(test_df['query_image']) | set(test_df['gallery_image']))} unique images, {len(test_df):,} pairs\")\nprint(f\"\\nStrategy: 5-Fold EVA02-Large + Alpha Masking + Focal Loss + Val mAP + Multi-Scale TTA + QE + Re-ranking\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:27:04.579079Z","iopub.execute_input":"2026-04-02T03:27:04.579783Z","iopub.status.idle":"2026-04-02T03:27:17.322228Z","shell.execute_reply.started":"2026-04-02T03:27:04.579761Z","shell.execute_reply":"2026-04-02T03:27:17.321318Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# JAGUAR RE-ID V8: ALL IMPROVEMENTS\n# ============================================================\n# Changes from V7:\n#   [1] Alpha mask support — zero out backgrounds to prevent\n#       spurious correlation with scenery\n#   [2] Focal Loss — directly addresses class imbalance for\n#       identity-balanced mAP metric\n#   [3] Val mAP monitoring + best checkpoint — stop flying blind\n#   [4] LR warmup — stabilize early fine-tuning of pretrained ViT\n#   [5] CLS + GeM concatenation — capture global + local features\n#   [6] Differential LR — lower for backbone, higher for head\n# ============================================================\n\nN_FOLDS = 5\nMODEL_NAME = 'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k'\nIMG_SIZE = 448\nBATCH_SIZE = 4\nGRAD_ACCUM = 4\nEPOCHS = 12\nBACKBONE_LR = 1e-5       # Lower LR for pretrained backbone\nHEAD_LR = 5e-5           # Higher LR for new layers (GeM, BN, ArcFace)\nWD = 1e-3\nWARMUP_EPOCHS = 2        # Linear warmup before cosine decay\nUSE_ALPHA_MASK = True    # Apply alpha mask to zero out backgrounds\nGRAD_CKPT = True         # Gradient checkpointing to reduce VRAM\n\n# Multi-scale TTA sizes (must be divisible by patch_size=14)\nTTA_SIZES = [392, 448, 518]\n\n# Focal loss params\nFOCAL_GAMMA = 2.0\nFOCAL_ALPHA = 0.25\n\nprint(f\"Plan: {N_FOLDS}-fold {MODEL_NAME.split('.')[0]}\")\nprint(f\"  {IMG_SIZE}px train | TTA at {TTA_SIZES}\")\nprint(f\"  {EPOCHS} epochs per fold | {WARMUP_EPOCHS} warmup epochs\")\nprint(f\"  Alpha masking: {USE_ALPHA_MASK}\")\nprint(f\"  Backbone LR: {BACKBONE_LR} | Head LR: {HEAD_LR}\")\nprint(f\"  Focal Loss: gamma={FOCAL_GAMMA}, alpha={FOCAL_ALPHA}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:27:17.322949Z","iopub.execute_input":"2026-04-02T03:27:17.323156Z","iopub.status.idle":"2026-04-02T03:27:17.328527Z","shell.execute_reply.started":"2026-04-02T03:27:17.323135Z","shell.execute_reply":"2026-04-02T03:27:17.32797Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# FOCAL LOSS — addresses class imbalance directly\n# ============================================================\nclass FocalLoss(nn.Module):\n    \"\"\"Focal Loss for handling class imbalance.\n    Down-weights easy examples so rare jaguars with fewer images\n    contribute more meaningfully to the loss.\"\"\"\n    def __init__(self, gamma=2.0, alpha=0.25, reduction='mean'):\n        super().__init__()\n        self.gamma = gamma\n        self.alpha = alpha\n        self.reduction = reduction\n\n    def forward(self, inputs, targets):\n        ce_loss = F.cross_entropy(inputs, targets, reduction='none')\n        pt = torch.exp(-ce_loss)\n        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss\n        if self.reduction == 'mean':\n            return focal_loss.mean()\n        return focal_loss\n\n\n# ============================================================\n# GeM POOLING\n# ============================================================\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\n# ============================================================\n# ARCFACE HEAD\n# ============================================================\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\n# ============================================================\n# MODEL — now with CLS + GeM concatenation\n# ============================================================\nclass JaguarModel(nn.Module):\n    def __init__(self, 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        if GRAD_CKPT:\n            self.backbone.set_grad_checkpointing(True)\n            print('  Gradient checkpointing enabled')\n        self.feat_dim = self.backbone.num_features\n        self.gem = GeM()\n\n        # CLS + GeM concat → project back to feat_dim\n        self.neck = nn.Sequential(\n            nn.Linear(self.feat_dim * 2, self.feat_dim),\n            nn.BatchNorm1d(self.feat_dim),\n        )\n        self.head = ArcFaceLayer(self.feat_dim, num_classes, s=s, m=m)\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            # CLS token is the first token\n            cls_token = features[:, 0, :]           # (B, C)\n            patch_tokens = features[:, 1:, :]       # (B, N-1, C)\n\n            H = W = int(math.sqrt(patch_tokens.shape[1]))\n            # Handle case where patch count isn't a perfect square\n            if H * W != patch_tokens.shape[1]:\n                H = W = int(math.sqrt(patch_tokens.shape[1]))\n                patch_tokens = patch_tokens[:, :H*W, :]\n\n            spatial = patch_tokens.permute(0, 2, 1).reshape(B, C, H, W)\n            gem_feat = self.gem(spatial).flatten(1)  # (B, C)\n\n            # Concatenate CLS + GeM for global + local features\n            emb = torch.cat([cls_token, gem_feat], dim=1)  # (B, 2C)\n        else:\n            # Fallback for CNN-style outputs\n            gem_feat = self.gem(features).flatten(1)\n            emb = torch.cat([gem_feat, gem_feat], dim=1)  # keep dimensions consistent\n\n        emb = self.neck(emb)  # (B, feat_dim)\n\n        if label is not None:\n            return self.head(emb, label)\n        return emb","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:27:17.330242Z","iopub.execute_input":"2026-04-02T03:27:17.330422Z","iopub.status.idle":"2026-04-02T03:27:17.38666Z","shell.execute_reply.started":"2026-04-02T03:27:17.330406Z","shell.execute_reply":"2026-04-02T03:27:17.385882Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# DATASET — with alpha mask support\n# ============================================================\nclass JaguarDataset(Dataset):\n    \"\"\"Loads images with optional alpha mask application.\n    When use_alpha=True, zeros out the background using the alpha channel,\n    forcing the model to learn from the jaguar's spot patterns only.\"\"\"\n\n    def __init__(self, df, img_dir, transform=None, is_test=False, use_alpha=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        self.use_alpha = use_alpha\n\n    def __len__(self):\n        return len(self.df)\n\n    def _load_image(self, filename):\n        path = os.path.join(self.img_dir, filename)\n        try:\n            img = Image.open(path)\n\n            if self.use_alpha and img.mode == 'RGBA':\n                # Extract alpha channel and apply as mask\n                r, g, b, a = img.split()\n                # Convert alpha to binary mask (0 or 255)\n                a_np = np.array(a).astype(np.float32) / 255.0\n                rgb_np = np.array(img.convert('RGB')).astype(np.float32)\n                # Apply mask: background becomes neutral gray (128)\n                # Using gray rather than black avoids creating artificial edges\n                masked = rgb_np * a_np[:, :, None] + 128.0 * (1 - a_np[:, :, None])\n                img = Image.fromarray(masked.astype(np.uint8))\n            else:\n                img = img.convert('RGB')\n\n            return img\n        except Exception:\n            return Image.new('RGB', (448, 448), (128, 128, 128))\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img = self._load_image(row['filename'])\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\n# ============================================================\n# TRANSFORMS — same as V7 (augmentations are adequate)\n# ============================================================\ndef get_train_transform():\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(size):\n    return T.Compose([\n        T.Resize((size, size)),\n        T.ToTensor(),\n        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n    ])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:27:17.387467Z","iopub.execute_input":"2026-04-02T03:27:17.38766Z","iopub.status.idle":"2026-04-02T03:27:17.406353Z","shell.execute_reply.started":"2026-04-02T03:27:17.387644Z","shell.execute_reply":"2026-04-02T03:27:17.405253Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# VALIDATION: Identity-Balanced mAP\n# ============================================================\ndef compute_val_map(model, val_df, img_dir, batch_size=4):\n    \"\"\"Compute identity-balanced mAP on the validation fold.\n    This mirrors the competition metric: mAP per identity, then averaged.\"\"\"\n    model.eval()\n\n    val_dataset = JaguarDataset(\n        val_df, img_dir, get_test_transform(IMG_SIZE),\n        is_test=False, use_alpha=USE_ALPHA_MASK\n    )\n    val_loader = DataLoader(val_dataset, batch_size=batch_size,\n                            shuffle=False, num_workers=2)\n\n    all_embs = []\n    all_labels = []\n    with torch.no_grad():\n        for imgs, labels in val_loader:\n            imgs = imgs.to(DEVICE)\n            with torch.amp.autocast('cuda'):\n                emb = model(imgs)\n            all_embs.append(emb.float().cpu())\n            all_labels.append(labels)\n\n    embs = torch.cat(all_embs, 0)\n    embs = F.normalize(embs, dim=1).numpy()\n    labels = torch.cat(all_labels, 0).numpy()\n\n    # Compute similarity matrix\n    sim = embs @ embs.T\n\n    # Identity-balanced mAP: compute AP per identity, then average\n    unique_labels = np.unique(labels)\n    per_identity_aps = []\n\n    for identity in unique_labels:\n        # Get indices for this identity\n        query_mask = labels == identity\n        query_indices = np.where(query_mask)[0]\n\n        if len(query_indices) < 2:\n            continue  # Need at least 2 images for query-gallery pairs\n\n        aps_for_identity = []\n        for q_idx in query_indices:\n            # Gallery = all images except this query\n            gallery_mask = np.ones(len(labels), dtype=bool)\n            gallery_mask[q_idx] = False\n            gallery_indices = np.where(gallery_mask)[0]\n\n            # Rank gallery by similarity to query\n            sims = sim[q_idx, gallery_indices]\n            sorted_idx = np.argsort(-sims)\n            sorted_labels = labels[gallery_indices[sorted_idx]]\n\n            # Compute AP for this query\n            relevant = (sorted_labels == identity).astype(float)\n            if relevant.sum() == 0:\n                continue\n            cumsum = np.cumsum(relevant)\n            precision_at_k = cumsum / np.arange(1, len(relevant) + 1)\n            ap = (precision_at_k * relevant).sum() / relevant.sum()\n            aps_for_identity.append(ap)\n\n        if aps_for_identity:\n            per_identity_aps.append(np.mean(aps_for_identity))\n\n    mean_ap = np.mean(per_identity_aps) if per_identity_aps else 0.0\n    return mean_ap","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:27:17.407427Z","iopub.execute_input":"2026-04-02T03:27:17.407683Z","iopub.status.idle":"2026-04-02T03:27:17.435486Z","shell.execute_reply.started":"2026-04-02T03:27:17.407658Z","shell.execute_reply":"2026-04-02T03:27:17.43433Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# LR SCHEDULER WITH WARMUP\n# ============================================================\nclass WarmupCosineScheduler:\n    \"\"\"Linear warmup followed by cosine decay.\n    Stabilizes early training when fine-tuning large pretrained ViTs.\"\"\"\n\n    def __init__(self, optimizer, warmup_epochs, total_epochs, steps_per_epoch):\n        self.optimizer = optimizer\n        self.warmup_steps = warmup_epochs * steps_per_epoch\n        self.total_steps = total_epochs * steps_per_epoch\n        self.base_lrs = [pg['lr'] for pg in optimizer.param_groups]\n        self.current_step = 0\n\n    def step(self):\n        self.current_step += 1\n        if self.current_step <= self.warmup_steps:\n            # Linear warmup\n            scale = self.current_step / max(1, self.warmup_steps)\n        else:\n            # Cosine decay\n            progress = (self.current_step - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps)\n            scale = 0.5 * (1 + math.cos(math.pi * progress))\n\n        for i, pg in enumerate(self.optimizer.param_groups):\n            pg['lr'] = self.base_lrs[i] * scale\n\n    def get_last_lr(self):\n        return [pg['lr'] for pg in self.optimizer.param_groups]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:27:17.436664Z","iopub.execute_input":"2026-04-02T03:27:17.437641Z","iopub.status.idle":"2026-04-02T03:27:17.46126Z","shell.execute_reply.started":"2026-04-02T03:27:17.437598Z","shell.execute_reply":"2026-04-02T03:27:17.460558Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# EXTRACT EMBEDDINGS WITH MULTI-SCALE TTA\n# ============================================================\ndef extract_multiscale_tta(model, unique_images, img_dir, batch_size=4):\n    \"\"\"Extract embeddings at multiple resolutions with horizontal flip.\n    For each image: len(TTA_SIZES) scales x 2 flips = 6 views averaged.\"\"\"\n    model.eval()\n    all_embs = []\n\n    for tta_size in TTA_SIZES:\n        loader = DataLoader(\n            JaguarDataset(\n                pd.DataFrame({'filename': unique_images}), img_dir,\n                get_test_transform(tta_size), is_test=True,\n                use_alpha=USE_ALPHA_MASK\n            ),\n            batch_size=batch_size, shuffle=False, num_workers=2\n        )\n        feats = []\n        with torch.no_grad():\n            for imgs, fnames in loader:\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 = (f1 + f2) / 2\n                feats.append(f_avg.float().cpu())\n        emb = torch.cat(feats, 0)\n        emb = F.normalize(emb, dim=1)\n        all_embs.append(emb)\n        del loader\n        print(f\"    TTA @ {tta_size}px done\")\n\n    # Average across scales and re-normalize\n    combined = torch.stack(all_embs).mean(dim=0)\n    combined = F.normalize(combined, dim=1)\n    return combined.numpy()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:27:17.462088Z","iopub.execute_input":"2026-04-02T03:27:17.46234Z","iopub.status.idle":"2026-04-02T03:27:17.486763Z","shell.execute_reply.started":"2026-04-02T03:27:17.462324Z","shell.execute_reply":"2026-04-02T03:27:17.485731Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# POST-PROCESSING\n# ============================================================\ndef query_expansion(emb, top_k=3):\n    print(\"  Applying Query Expansion...\")\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))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:27:17.487846Z","iopub.execute_input":"2026-04-02T03:27:17.488086Z","iopub.status.idle":"2026-04-02T03:27:17.512737Z","shell.execute_reply.started":"2026-04-02T03:27:17.48807Z","shell.execute_reply":"2026-04-02T03:27:17.512087Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# MAIN: 5-FOLD TRAINING WITH ALL IMPROVEMENTS\n# ============================================================\nunique_test = sorted(set(test_df['query_image']) | set(test_df['gallery_image']))\nimg_map = {n: i for i, n in enumerate(unique_test)}\n\nprint(f\"\\nUnique test images: {len(unique_test)}\")\nprint(f\"\\n{'#'*60}\")\nprint(f\"# V8: 5-FOLD EVA02-Large + Alpha Mask + Focal Loss\")\nprint(f\"# + Val mAP + LR Warmup + CLS+GeM + Multi-Scale TTA\")\nprint(f\"# TTA: {len(TTA_SIZES)} scales x 2 flips = {len(TTA_SIZES)*2} views per image\")\nprint(f\"# Then: QE + K-Reciprocal Re-ranking\")\nprint(f\"{'#'*60}\")\n\nskf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42)\n\nall_embeddings = []  # Store normalized embeddings from each fold\nfold_val_maps = []   # Track best val mAP per fold\n\nfor fold, (train_idx, val_idx) in enumerate(skf.split(train_df, train_df['label'])):\n    seed_everything(42 + fold)\n\n    print(f\"\\n{'='*60}\")\n    print(f\"FOLD {fold+1}/{N_FOLDS} | Train: {len(train_idx)} | Val: {len(val_idx)}\")\n    print(f\"{'='*60}\")\n\n    fold_train = train_df.iloc[train_idx].copy()\n    fold_val = train_df.iloc[val_idx].copy()\n\n    model = JaguarModel(NUM_CLASSES).to(DEVICE)\n    print(f\"  Params: {sum(p.numel() for p in model.parameters()):,}\")\n\n    train_loader = DataLoader(\n        JaguarDataset(fold_train, TRAIN_DIR, get_train_transform(),\n                      use_alpha=USE_ALPHA_MASK),\n        batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True\n    )\n\n    # ---- Differential LR: lower for backbone, higher for head ----\n    backbone_params = list(model.backbone.parameters())\n    head_params = (list(model.gem.parameters()) +\n                   list(model.neck.parameters()) +\n                   list(model.head.parameters()))\n\n    optimizer = torch.optim.AdamW([\n        {'params': backbone_params, 'lr': BACKBONE_LR},\n        {'params': head_params, 'lr': HEAD_LR},\n    ], weight_decay=WD)\n\n    # Steps-level warmup + cosine scheduler\n    steps_per_epoch = math.ceil(len(train_loader) / GRAD_ACCUM)\n    scheduler = WarmupCosineScheduler(optimizer, WARMUP_EPOCHS, EPOCHS, steps_per_epoch)\n\n    scaler = torch.amp.GradScaler('cuda')\n    criterion = FocalLoss(gamma=FOCAL_GAMMA, alpha=FOCAL_ALPHA)\n\n    best_val_map = 0.0\n    ckpt_path = f'/kaggle/working/best_fold{fold}.pt'\n\n    # ---- Train this fold ----\n    for epoch in range(EPOCHS):\n        model.train()\n        total_loss = 0\n        optimizer.zero_grad()\n\n        for i, (imgs, labels) in enumerate(tqdm(train_loader, leave=False, desc=f\"F{fold+1} E{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\n            if (i + 1) % GRAD_ACCUM == 0:\n                scaler.step(optimizer)\n                scaler.update()\n                optimizer.zero_grad()\n                scheduler.step()  # Step-level scheduling\n\n            total_loss += loss.item() * GRAD_ACCUM\n\n        # Handle remaining gradients\n        scaler.step(optimizer)\n        scaler.update()\n        optimizer.zero_grad()\n\n        avg_loss = total_loss / len(train_loader)\n        lrs = scheduler.get_last_lr()\n\n        # ---- Validate every 2 epochs or on the last epoch ----\n        if (epoch + 1) % 2 == 0 or epoch == EPOCHS - 1:\n            gc.collect()\n            torch.cuda.empty_cache()\n            val_map = compute_val_map(model, fold_val, TRAIN_DIR, batch_size=BATCH_SIZE)\n            improved = ''\n            if val_map > best_val_map:\n                best_val_map = val_map\n                torch.save(model.state_dict(), ckpt_path)\n                improved = ' ★ BEST'\n            print(f\"  Epoch {epoch+1}/{EPOCHS} | Loss: {avg_loss:.4f} | \"\n                  f\"BB_LR: {lrs[0]:.2e} | Head_LR: {lrs[1]:.2e} | \"\n                  f\"Val mAP: {val_map:.4f}{improved}\")\n        else:\n            print(f\"  Epoch {epoch+1}/{EPOCHS} | Loss: {avg_loss:.4f} | \"\n                  f\"BB_LR: {lrs[0]:.2e} | Head_LR: {lrs[1]:.2e}\")\n\n    # ---- Load best checkpoint for this fold ----\n    print(f\"  Loading best checkpoint (val mAP: {best_val_map:.4f})\")\n    model.load_state_dict(torch.load(ckpt_path, map_location=DEVICE))\n    fold_val_maps.append(best_val_map)\n\n    # ---- Extract embeddings with multi-scale TTA ----\n    print(f\"  Extracting multi-scale TTA embeddings...\")\n    fold_emb = extract_multiscale_tta(model, unique_test, TEST_DIR, batch_size=BATCH_SIZE)\n    all_embeddings.append(fold_emb)\n    print(f\"  Fold {fold+1} done! Embeddings: {fold_emb.shape} | Best val mAP: {best_val_map:.4f}\")\n\n    # Free memory\n    del model, optimizer, scaler, scheduler, train_loader\n    gc.collect()\n    torch.cuda.empty_cache()\n\nprint(f\"\\nPer-fold val mAP: {[f'{m:.4f}' for m in fold_val_maps]}\")\nprint(f\"Mean val mAP: {np.mean(fold_val_maps):.4f} ± {np.std(fold_val_maps):.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:27:17.513603Z","iopub.execute_input":"2026-04-02T03:27:17.513778Z","iopub.status.idle":"2026-04-02T03:27:32.228773Z","shell.execute_reply.started":"2026-04-02T03:27:17.513762Z","shell.execute_reply":"2026-04-02T03:27:32.22729Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# ENSEMBLE ALL FOLDS + POST-PROCESSING\n# ============================================================\nprint(f\"\\n{'='*60}\")\nprint(f\"ENSEMBLING {len(all_embeddings)} FOLDS\")\nprint(f\"{'='*60}\")\n\n# Average embeddings across folds, then normalize\navg_emb = np.mean(all_embeddings, axis=0)\navg_emb = avg_emb / np.linalg.norm(avg_emb, axis=1, keepdims=True)\n\n# Query Expansion on the averaged embeddings\navg_emb = query_expansion(avg_emb)\n\n# Compute similarity matrix\nsim_matrix = avg_emb @ avg_emb.T\n\n# Re-ranking\nsim_matrix = k_reciprocal_rerank(sim_matrix)\n\n# ---- Generate Submission ----\nprint(\"\\nGenerating submission...\")\npreds = []\nfor _, row in tqdm(test_df.iterrows(), total=len(test_df), desc=\"Mapping\"):\n    s = sim_matrix[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! V8 Complete\")\nprint(f\"{'='*60}\")\nprint(f\"Folds: {len(all_embeddings)}\")\nprint(f\"Per-fold val mAP: {[f'{m:.4f}' for m in fold_val_maps]}\")\nprint(f\"Mean val mAP: {np.mean(fold_val_maps):.4f} ± {np.std(fold_val_maps):.4f}\")\nprint(f\"TTA scales: {TTA_SIZES} x 2 flips = {len(TTA_SIZES)*2} views\")\nprint(f\"Alpha masking: {USE_ALPHA_MASK}\")\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))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-02T03:27:32.230551Z","iopub.status.idle":"2026-04-02T03:27:32.230926Z","shell.execute_reply.started":"2026-04-02T03:27:32.230746Z","shell.execute_reply":"2026-04-02T03:27:32.230763Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}