{"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":[{"sourceType":"competition","sourceId":126777,"databundleVersionId":15314950,"isSourceIdPinned":false}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nv03: ArcFace Metric Learning\n==============================\nTrain EfficientNet-B0 with ArcFace loss — the standard approach for\nclosed-set re-identification tasks.\n\nKey differences from v02 (classifier):\n  Loss:    ArcFace (angular margin m=0.5, scale s=30) vs CrossEntropy\n  Embedding: 512-dim projected (Linear + BN) vs 1280-dim raw\n  Normalise: embeddings L2-normalised DURING training (required by ArcFace)\n  LR:      backbone lr=1e-5 (lower than v02's 1e-4) — ArcFace gradients larger\n  Monitor: val mAP every VAL_MAP_INTERVAL epochs; best checkpoint by mAP\n  Clip:    gradient norm clipped at 1.0 for ArcFace stability\n\nStrategy:\n  Stage 1 (5 epochs)  — frozen backbone, train projection + ArcFace head\n  Stage 2 (30 epochs) — full fine-tune, differential LR\n  Then retrain on ALL data for final submission.\n\nRun on Kaggle (T4 or P100 GPU):\n  Attach the jaguar-re-id competition dataset.\n  Outputs land in /kaggle/working/output/ — download the zip.\n\"\"\"\n\nimport sys, os, time\n\n# -- Clone / update repo -------------------------------------------------------\nREPO_URL = 'https://github.com/Smooth-Cactus0/jaguar-re-identification.git'\nREPO_DIR = '/kaggle/working/repo'\n\nif os.path.exists(REPO_DIR):\n    os.system(f'git -C {REPO_DIR} pull --ff-only')\nelse:\n    os.system(f'git clone {REPO_URL} {REPO_DIR}')\nsys.path.insert(0, REPO_DIR)\n\nimport math\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import OneCycleLR\nfrom pathlib import Path\nimport warnings\nwarnings.filterwarnings('ignore')\n\nfrom src.data      import JaguarDataset, get_transforms, encode_labels, get_identity_splits\nfrom src.models    import EmbeddingModel\nfrom src.losses    import ArcFaceLoss\nfrom src.inference import extract_embeddings, make_submission\nfrom src.evaluate  import compute_map, print_results, save_benchmark\n\n# -- Paths (Kaggle layout) -----------------------------------------------------\nKAGGLE_INPUT = Path('/kaggle/input/competitions/jaguar-re-id')\nTRAIN_DIR    = KAGGLE_INPUT / 'train' / 'train'\nTEST_DIR     = KAGGLE_INPUT / 'test'  / 'test'\nOUT_DIR      = Path('/kaggle/working/output')\nOUT_DIR.mkdir(exist_ok=True)\n\n# -- Config --------------------------------------------------------------------\nBACKBONE        = 'efficientnet_b0'\nEMBEDDING_DIM   = 512          # projected dim (1280 -> 512 via Linear + BN)\nIMG_SIZE        = 224\nBATCH_SIZE      = 64\nN_CLASSES       = 31\nLR_HEAD         = 1e-3         # Stage 1 and ArcFace head in Stage 2\nLR_BACKBONE     = 1e-5         # lower than v02 — ArcFace gradients are stronger\nEPOCHS_S1       = 5\nEPOCHS_S2       = 30\nWEIGHT_DECAY    = 1e-4\nARC_MARGIN      = 0.5          # angular margin in radians (~28.6°)\nARC_SCALE       = 30.0         # feature scale; lower for small-scale re-ID\nGRAD_CLIP_NORM  = 1.0          # gradient clipping for ArcFace stability\nVAL_MAP_INTERVAL = 5           # compute val mAP every N epochs of stage 2\nDEVICE          = 'cuda' if torch.cuda.is_available() else 'cpu'\nVERSION         = 'v03'\nCV_FOLD         = 0\n\nprint('\\n' + '='*60)\nprint(f'  {VERSION}: ArcFace Metric Learning')\nprint(f'  Backbone      : {BACKBONE}')\nprint(f'  Embedding dim : {EMBEDDING_DIM}')\nprint(f'  ArcFace       : margin={ARC_MARGIN}, scale={ARC_SCALE}')\nprint(f'  Device        : {DEVICE}')\nprint(f'  Epochs        : S1={EPOCHS_S1} + S2={EPOCHS_S2}')\nprint(f'  Batch size    : {BATCH_SIZE}')\nprint('='*60)\n\n# -- Load data -----------------------------------------------------------------\nprint('\\n[1/8] Loading data...')\ntrain_df = pd.read_csv(KAGGLE_INPUT / 'train.csv')\ntest_df  = pd.read_csv(KAGGLE_INPUT / 'test.csv')\nlabels, label_to_idx, idx_to_label = encode_labels(train_df['ground_truth'])\nprint(f'  Train: {len(train_df)} images, {N_CLASSES} identities')\n\n# -- CV split ------------------------------------------------------------------\nprint(f'\\n[2/8] Creating CV split (fold {CV_FOLD}/5)...')\nfor fold, train_idx, val_idx in get_identity_splits(train_df, n_splits=5):\n    if fold == CV_FOLD:\n        break\n\ntf_train = get_transforms(mode='train', img_size=IMG_SIZE)\ntf_val   = get_transforms(mode='val',   img_size=IMG_SIZE)\n\ntrain_ds = JaguarDataset(\n    filenames=train_df['filename'].iloc[train_idx].tolist(),\n    img_dir=TRAIN_DIR, labels=labels[train_idx], transform=tf_train)\nval_ds = JaguarDataset(\n    filenames=train_df['filename'].iloc[val_idx].tolist(),\n    img_dir=TRAIN_DIR, labels=labels[val_idx], transform=tf_val)\n\ntrain_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,\n                          num_workers=2, pin_memory=True)\nval_loader   = DataLoader(val_ds,   batch_size=BATCH_SIZE, shuffle=False,\n                          num_workers=2, pin_memory=True)\nprint(f'  Train: {len(train_ds)} imgs | Val: {len(val_ds)} imgs')\n\n\n# -- Model architecture --------------------------------------------------------\nprint('\\n[3/8] Building model...')\n\n\nclass ArcFaceModel(nn.Module):\n    \"\"\"\n    EfficientNet backbone + projection head for ArcFace training.\n\n    Forward always returns L2-normalised embeddings — ArcFaceLoss takes\n    them directly. Same interface is used for inference (extract_embeddings).\n\n    Projection: Linear(native_dim, embedding_dim, bias=False) + BatchNorm1d\n    bias=False because BN provides its own learnable offset.\n    \"\"\"\n\n    def __init__(self, backbone_name: str, embedding_dim: int,\n                 pretrained: bool = True):\n        super().__init__()\n        self.backbone  = EmbeddingModel(backbone_name, pretrained=pretrained)\n        native_dim     = self.backbone.out_dim\n        self.proj      = nn.Sequential(\n            nn.Linear(native_dim, embedding_dim, bias=False),\n            nn.BatchNorm1d(embedding_dim),\n        )\n        self.out_dim = embedding_dim\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        feat = self.backbone(x)                          # (B, native_dim)\n        emb  = self.proj(feat)                           # (B, embedding_dim)\n        return F.normalize(emb, p=2, dim=1)              # (B, embedding_dim), unit norm\n\n    def freeze_backbone(self):\n        for p in self.backbone.parameters():\n            p.requires_grad_(False)\n\n    def unfreeze_backbone(self):\n        for p in self.backbone.parameters():\n            p.requires_grad_(True)\n\n\nmodel   = ArcFaceModel(BACKBONE, EMBEDDING_DIM).to(DEVICE)\narcface = ArcFaceLoss(EMBEDDING_DIM, N_CLASSES, margin=ARC_MARGIN, scale=ARC_SCALE).to(DEVICE)\n\ntotal_backbone = sum(p.numel() for p in model.backbone.parameters())\ntotal_proj     = sum(p.numel() for p in model.proj.parameters())\ntotal_arc      = sum(p.numel() for p in arcface.parameters())\nprint(f'  Backbone params : {total_backbone:,}')\nprint(f'  Projection params: {total_proj:,}')\nprint(f'  ArcFace params  : {total_arc:,}')\n\n\n# -- Training utilities --------------------------------------------------------\n\ndef run_epoch(model, arcface_loss, loader, optimizer, scheduler, is_train, device):\n    \"\"\"Run one epoch; return average ArcFace loss.\"\"\"\n    model.train(is_train)\n    arcface_loss.train(is_train)\n    total_loss, n = 0.0, 0\n    ctx = torch.enable_grad() if is_train else torch.no_grad()\n    with ctx:\n        for imgs, targets in loader:\n            imgs, targets = imgs.to(device), targets.to(device)\n            if is_train:\n                optimizer.zero_grad()\n            embs = model(imgs)                # (B, 512), L2-normalised\n            loss = arcface_loss(embs, targets)\n            if is_train:\n                loss.backward()\n                # Clip gradients — ArcFace margin creates sharp spikes early in training\n                torch.nn.utils.clip_grad_norm_(\n                    list(model.parameters()) + list(arcface_loss.parameters()),\n                    max_norm=GRAD_CLIP_NORM)\n                optimizer.step()\n                scheduler.step()\n            total_loss += loss.item() * len(imgs)\n            n          += len(imgs)\n    return total_loss / n\n\n\ndef compute_val_map(model, val_ds, val_labels, batch_size, device):\n    \"\"\"Extract val embeddings and return identity-balanced mAP.\"\"\"\n    embs    = extract_embeddings(model, val_ds, batch_size=batch_size,\n                                 device=device, l2_normalise=False,  # already normalised\n                                 desc='Val mAP')\n    results = compute_map(embs, val_labels, identity_balanced=True)\n    return results\n\n\n# -- Stage 1: head-only training -----------------------------------------------\nprint('\\n[4/8] Stage 1: projection + ArcFace head only (backbone frozen)...')\nmodel.freeze_backbone()\ntrainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\nprint(f'  Trainable params: {trainable:,} + {total_arc:,} (ArcFace) = {trainable+total_arc:,}')\n\n# Include ArcFace weight matrix in optimiser\nhead_params = list(filter(lambda p: p.requires_grad, model.parameters())) + \\\n              list(arcface.parameters())\nopt_s1 = AdamW(head_params, lr=LR_HEAD, weight_decay=WEIGHT_DECAY)\nsch_s1 = OneCycleLR(opt_s1, max_lr=LR_HEAD,\n                    steps_per_epoch=len(train_loader), epochs=EPOCHS_S1)\n\nhistory = {'train_loss': [], 'val_loss': []}\nfor ep in range(1, EPOCHS_S1 + 1):\n    t0 = time.time()\n    tr_loss = run_epoch(model, arcface, train_loader, opt_s1, sch_s1, True,  DEVICE)\n    va_loss = run_epoch(model, arcface, val_loader,   opt_s1, sch_s1, False, DEVICE)\n    history['train_loss'].append(tr_loss)\n    history['val_loss'].append(va_loss)\n    print(f'  S1 ep {ep:02d}/{EPOCHS_S1} | '\n          f'loss {tr_loss:.4f}/{va_loss:.4f} | '\n          f'{time.time()-t0:.0f}s')\n\n\n# -- Stage 2: full fine-tuning -------------------------------------------------\nprint('\\n[5/8] Stage 2: full fine-tuning...')\nmodel.unfreeze_backbone()\ntrainable_all = sum(p.numel() for p in model.parameters() if p.requires_grad)\nprint(f'  Trainable params: {trainable_all:,} + {total_arc:,} (ArcFace)')\n\nparam_groups = [\n    {'params': model.backbone.parameters(), 'lr': LR_BACKBONE},\n    {'params': model.proj.parameters(),     'lr': LR_HEAD},\n    {'params': arcface.parameters(),        'lr': LR_HEAD},\n]\nopt_s2 = AdamW(param_groups, weight_decay=WEIGHT_DECAY)\nsch_s2 = OneCycleLR(opt_s2,\n                    max_lr=[LR_BACKBONE, LR_HEAD, LR_HEAD],\n                    steps_per_epoch=len(train_loader),\n                    epochs=EPOCHS_S2)\n\nbest_val_map  = 0.0\nbest_state    = None\nval_map_log   = []  # (epoch, val_mAP)\n\nfor ep in range(1, EPOCHS_S2 + 1):\n    t0 = time.time()\n    tr_loss = run_epoch(model, arcface, train_loader, opt_s2, sch_s2, True,  DEVICE)\n    va_loss = run_epoch(model, arcface, val_loader,   opt_s2, sch_s2, False, DEVICE)\n    history['train_loss'].append(tr_loss)\n    history['val_loss'].append(va_loss)\n\n    # Evaluate val mAP every VAL_MAP_INTERVAL epochs\n    if ep % VAL_MAP_INTERVAL == 0 or ep == EPOCHS_S2:\n        val_results = compute_val_map(model, val_ds, labels[val_idx],\n                                      BATCH_SIZE, DEVICE)\n        vm = val_results['map']\n        val_map_log.append((EPOCHS_S1 + ep, vm))\n\n        marker = ''\n        if vm > best_val_map:\n            best_val_map = vm\n            best_state   = {k: v.cpu().clone() for k, v in model.state_dict().items()}\n            marker = '  ** best mAP **'\n        print(f'  S2 ep {ep:02d}/{EPOCHS_S2} | '\n              f'loss {tr_loss:.4f}/{va_loss:.4f} | '\n              f'val_mAP {vm:.4f} | '\n              f'{time.time()-t0:.0f}s{marker}')\n    else:\n        print(f'  S2 ep {ep:02d}/{EPOCHS_S2} | '\n              f'loss {tr_loss:.4f}/{va_loss:.4f} | '\n              f'{time.time()-t0:.0f}s')\n\nmodel.load_state_dict(best_state)\nprint(f'  Loaded best checkpoint (val_mAP={best_val_map:.4f})')\n\n\n# -- Evaluate on validation fold -----------------------------------------------\nprint('\\n[6/8] Computing final identity-balanced mAP...')\nval_results = compute_val_map(model, val_ds, labels[val_idx], BATCH_SIZE, DEVICE)\nprint_results(val_results, idx_to_label, title=f'{VERSION} - Fold {CV_FOLD} Validation')\n\n# Full train leave-one-out mAP for comparison with prior versions\nfull_ds = JaguarDataset(\n    filenames=train_df['filename'].tolist(),\n    img_dir=TRAIN_DIR, labels=labels, transform=tf_val)\nfull_embs    = extract_embeddings(model, full_ds, batch_size=BATCH_SIZE,\n                                   device=DEVICE, l2_normalise=False,\n                                   desc='Full train embeddings')\nfull_results = compute_map(full_embs, labels, identity_balanced=True)\nprint(f'\\n  Full train leave-one-out mAP: {full_results[\"map\"]:.4f}')\nprint(f'  v01: 0.2781  |  v02: 0.5258  |  v03: {full_results[\"map\"]:.4f}')\n\nsave_benchmark(\n    OUT_DIR / 'benchmarks_v03.csv',\n    {\n        'version':        VERSION,\n        'backbone':       BACKBONE,\n        'loss':           f'ArcFace (m={ARC_MARGIN}, s={ARC_SCALE})',\n        'embedding_dim':  EMBEDDING_DIM,\n        'img_size':       IMG_SIZE,\n        'augmentation':   'none (v05 adds augmentation)',\n        'cv_map_val':     round(val_results['map'],  5),\n        'cv_map_full':    round(full_results['map'], 5),\n        'best_identity':  idx_to_label[max(full_results['per_identity_ap'],\n                                           key=full_results['per_identity_ap'].get)],\n        'worst_identity': idx_to_label[min(full_results['per_identity_ap'],\n                                           key=full_results['per_identity_ap'].get)],\n        'epochs_s1':      EPOCHS_S1,\n        'epochs_s2':      EPOCHS_S2,\n        'notes':          'ArcFace, proj 512-dim BN, differential LR, best by val mAP',\n    }\n)\n\n\n# -- Visualisations ------------------------------------------------------------\nprint('\\n[7/8] Generating visualisations...')\n\n# 7a. Training loss curves + val mAP overlay\nep_range   = list(range(1, EPOCHS_S1 + EPOCHS_S2 + 1))\nmap_epochs = [x[0] for x in val_map_log]\nmap_vals   = [x[1] for x in val_map_log]\n\nfig, axes = plt.subplots(1, 2, figsize=(14, 4))\n\naxes[0].plot(ep_range, history['train_loss'], label='Train', color='steelblue')\naxes[0].plot(ep_range, history['val_loss'],   label='Val',   color='coral')\naxes[0].axvline(EPOCHS_S1, color='gray', linestyle='--', linewidth=1,\n                label=f'Unfreeze at ep {EPOCHS_S1}')\naxes[0].set_xlabel('Epoch'); axes[0].set_ylabel('ArcFace Loss')\naxes[0].set_title('Training Loss', fontweight='bold')\naxes[0].legend()\n\nax2 = axes[1]\nax2.plot(map_epochs, map_vals, 'o-', color='steelblue', markersize=6, label='Val mAP (v03)')\nax2.axhline(0.4821, color='coral',  linestyle='--', linewidth=1, label='v02 val mAP 0.4821')\nax2.axhline(0.2781, color='gray',   linestyle=':',  linewidth=1, label='v01 mAP 0.2781')\nax2.set_xlabel('Epoch'); ax2.set_ylabel('Identity-balanced mAP')\nax2.set_title('Validation mAP', fontweight='bold')\nax2.set_ylim(0, 1); ax2.legend()\n\nplt.suptitle(f'v03: ArcFace {BACKBONE} — Training Curves', fontsize=12, fontweight='bold')\nplt.tight_layout()\nfig.savefig(OUT_DIR / f'{VERSION}_training_curves.png', dpi=150, bbox_inches='tight')\nplt.close()\nprint('  Saved -> v03_training_curves.png')\n\n# 7b. Per-identity AP bar chart: v01 vs v02 vs v03\nv03_per_id = {idx_to_label[k]: v for k, v in full_results['per_identity_ap'].items()}\nv01_per_id = {\n    'Katniss':0.5694,'Lua':0.5186,'Bagua':0.5112,'Tomas':0.4515,\n    'Kamaikua':0.3999,'Pyte':0.3825,'Benita':0.3686,'Pollyanna':0.3640,\n    'Guaraci':0.3479,'Madalena':0.3408,'Pixana':0.3251,'Estella':0.3035,\n    'Abril':0.2865,'Apeiara':0.2752,'Ariely':0.2585,'Ti':0.2456,\n    'Alira':0.2420,'Akaloi':0.2377,'Solar':0.2270,'Ousado':0.2187,\n    'Medrosa':0.2149,'Overa':0.2058,'Kwang':0.1977,'Marcela':0.1883,\n    'Saseka':0.1687,'Oxum':0.1655,'Bororo':0.1604,'Bernard':0.1352,\n    'Jaju':0.1332,'Ipepo':0.0968,'Patricia':0.0797,\n}\n\nids_sorted = sorted(v03_per_id.keys(), key=lambda x: -v03_per_id[x])\nx = np.arange(len(ids_sorted))\nw = 0.28\n\nfig, ax = plt.subplots(figsize=(16, 5))\nax.bar(x - w, [v01_per_id.get(i, 0) for i in ids_sorted], w,\n       color='lightsteelblue', label='v01 frozen', edgecolor='white')\nax.bar(x,     [v03_per_id[i] for i in ids_sorted], w,\n       color='steelblue', label='v03 ArcFace', edgecolor='white')\nax.axhline(full_results['map'], color='navy', linestyle='--', linewidth=1,\n           label=f'v03 mAP={full_results[\"map\"]:.4f}')\nax.axhline(0.2781, color='gray', linestyle=':', linewidth=1, label='v01 mAP=0.2781')\nax.set_xticks(x); ax.set_xticklabels(ids_sorted, rotation=45, ha='right', fontsize=8)\nax.set_ylabel('Average Precision'); ax.set_ylim(0, 1)\nax.set_title('v01 vs v03 ArcFace: Per-Identity AP', fontsize=12, fontweight='bold')\nax.legend(fontsize=9)\nplt.tight_layout()\nfig.savefig(OUT_DIR / f'{VERSION}_per_identity_ap.png', dpi=150, bbox_inches='tight')\nplt.close()\nprint('  Saved -> v03_per_identity_ap.png')\n\n# 7c. t-SNE of ArcFace embeddings\nprint('  Computing t-SNE...')\nfrom sklearn.manifold import TSNE\nimport seaborn as sns\n\ntsne   = TSNE(n_components=2, perplexity=30, n_iter=1000, random_state=42, n_jobs=-1)\nembs2d = tsne.fit_transform(full_embs)\npalette = sns.color_palette('tab20', N_CLASSES) + \\\n          sns.color_palette('tab20b', max(0, N_CLASSES - 20))\n\nfig, ax = plt.subplots(figsize=(12, 10))\nfor lbl in range(N_CLASSES):\n    mask = labels == lbl\n    ax.scatter(embs2d[mask, 0], embs2d[mask, 1], color=palette[lbl],\n               label=idx_to_label[lbl], s=20, alpha=0.7, edgecolors='none')\nax.set_title(f'v03: t-SNE of ArcFace {BACKBONE} Embeddings (512-dim)',\n             fontsize=12, fontweight='bold')\nax.legend(loc='upper right', fontsize=6, ncol=2, markerscale=1.5,\n          framealpha=0.7, title='Identity')\nax.axis('off')\nplt.tight_layout()\nfig.savefig(OUT_DIR / f'{VERSION}_tsne.png', dpi=150, bbox_inches='tight')\nplt.close()\nprint('  Saved -> v03_tsne.png')\n\n\n# -- Full retrain + submission -------------------------------------------------\nprint('\\n[8/8] Retraining on ALL data for submission...')\nfull_train_ds = JaguarDataset(\n    filenames=train_df['filename'].tolist(),\n    img_dir=TRAIN_DIR, labels=labels, transform=tf_train)\nfull_loader = DataLoader(full_train_ds, batch_size=BATCH_SIZE, shuffle=True,\n                         num_workers=2, pin_memory=True)\n\nmodel_final   = ArcFaceModel(BACKBONE, EMBEDDING_DIM).to(DEVICE)\narcface_final = ArcFaceLoss(EMBEDDING_DIM, N_CLASSES,\n                             margin=ARC_MARGIN, scale=ARC_SCALE).to(DEVICE)\n\n# Stage 1: head only\nmodel_final.freeze_backbone()\nhead_params_f = list(filter(lambda p: p.requires_grad, model_final.parameters())) + \\\n                list(arcface_final.parameters())\nopt_f1 = AdamW(head_params_f, lr=LR_HEAD, weight_decay=WEIGHT_DECAY)\nsch_f1 = OneCycleLR(opt_f1, max_lr=LR_HEAD,\n                    steps_per_epoch=len(full_loader), epochs=EPOCHS_S1)\nfor ep in range(1, EPOCHS_S1 + 1):\n    loss = run_epoch(model_final, arcface_final, full_loader, opt_f1, sch_f1, True, DEVICE)\n    print(f'  Full S1 ep {ep}/{EPOCHS_S1} | loss {loss:.4f}')\n\n# Stage 2: full fine-tune\nmodel_final.unfreeze_backbone()\nparam_groups_f = [\n    {'params': model_final.backbone.parameters(), 'lr': LR_BACKBONE},\n    {'params': model_final.proj.parameters(),     'lr': LR_HEAD},\n    {'params': arcface_final.parameters(),        'lr': LR_HEAD},\n]\nopt_f2 = AdamW(param_groups_f, weight_decay=WEIGHT_DECAY)\nsch_f2 = OneCycleLR(opt_f2,\n                    max_lr=[LR_BACKBONE, LR_HEAD, LR_HEAD],\n                    steps_per_epoch=len(full_loader),\n                    epochs=EPOCHS_S2)\nfor ep in range(1, EPOCHS_S2 + 1):\n    loss = run_epoch(model_final, arcface_final, full_loader, opt_f2, sch_f2, True, DEVICE)\n    if ep % 5 == 0:\n        print(f'  Full S2 ep {ep}/{EPOCHS_S2} | loss {loss:.4f}')\n\n# Test embeddings + submission\ntest_filenames = sorted(pd.unique(\n    test_df[['query_image', 'gallery_image']].values.ravel()))\ntest_ds = JaguarDataset(\n    filenames=test_filenames, img_dir=TEST_DIR, labels=None, transform=tf_val)\nmodel_final.eval()\ntest_embs = extract_embeddings(model_final, test_ds, batch_size=BATCH_SIZE,\n                                device=DEVICE, l2_normalise=False,\n                                desc='Test embeddings')\nfname_to_idx = {f: i for i, f in enumerate(test_filenames)}\nmake_submission(\n    test_df=test_df, test_embeddings=test_embs,\n    filename_to_idx=fname_to_idx,\n    output_path=OUT_DIR / f'submission_{VERSION}.csv')\n\ntorch.save(model_final.state_dict(), OUT_DIR / f'model_{VERSION}.pth')\nprint(f'  Model saved -> model_{VERSION}.pth')\n\nprint(f'\\n{\"=\"*60}')\nprint(f'  {VERSION} COMPLETE')\nprint(f'  Val fold mAP  : {val_results[\"map\"]:.4f}')\nprint(f'  Full train mAP: {full_results[\"map\"]:.4f}')\nprint(f'  Best val mAP  : {best_val_map:.4f}')\nprint(f'  Outputs in    : {OUT_DIR}')\nprint(f'{\"=\"*60}')\nprint('\\nDownload /kaggle/working/output/ zip and push figures + results to GitHub.')","metadata":{"_uuid":"d11d29ac-55c9-43ea-8bd1-e73423f6cf8a","_cell_guid":"b21d4c9f-fb36-4e57-ab64-9dd3c1ed9018","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2026-03-03T18:07:14.008138Z","iopub.execute_input":"2026-03-03T18:07:14.008789Z","execution_failed":"2026-03-03T18:08:13.172Z"}},"outputs":[],"execution_count":null}]}