{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":126777,"databundleVersionId":15314950,"isSourceIdPinned":false}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nv05: Augmentation + Test-Time Augmentation\n==========================================\nBest backbone from v04: ConvNeXt-Small (full_mAP=0.8382, LB=0.797).\n\nTwo improvements over v04:\n  1. TRAINING AUGMENTATION — use augment_level='heavy' during training:\n       RandomRotation(10°), aggressive ColorJitter, RandomErasing(p=0.25)\n     Forces the model to learn robust features rather than memorising\n     exact viewpoint/lighting conditions.\n\n  2. TEST-TIME AUGMENTATION (TTA) — at inference, run each image through\n     N augmented views and average the L2-normalised embeddings:\n       - original centre crop\n       - horizontal flip\n       - top-left crop, top-right crop, bottom-centre crop\n     Averaging unit-norm vectors on the hypersphere is equivalent to\n     finding the centroid direction — more robust than a single view.\n\nBoth changes require no architecture modification; same\nConvNeXt-Small + ArcFace(m=0.5, s=30) + 512-dim projection.\nCaching + AMP retained for speed.\n\"\"\"\n\nimport sys, os, time\n\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 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 torch.cuda.amp import GradScaler, autocast\nfrom torchvision import transforms\nfrom pathlib import Path\nimport warnings\nwarnings.filterwarnings('ignore')\n\nfrom src.data      import (JaguarDataset, get_transforms, encode_labels,\n                            get_identity_splits, build_image_cache,\n                            IMAGENET_MEAN, IMAGENET_STD)\nfrom src.models    import EmbeddingModel\nfrom src.losses    import ArcFaceLoss\nfrom src.inference import make_submission\nfrom src.evaluate  import compute_map, print_results, save_benchmark\n\n# -- Paths ---------------------------------------------------------------------\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       = 'convnext_small'   # winner from v04\nEMBEDDING_DIM  = 512\nIMG_SIZE       = 224\nBATCH_SIZE     = 64\nN_CLASSES      = 31\nLR_HEAD        = 1e-3\nLR_BACKBONE    = 1e-5\nEPOCHS_S1      = 5\nEPOCHS_S2      = 25\nWEIGHT_DECAY   = 1e-4\nARC_MARGIN     = 0.5\nARC_SCALE      = 30.0\nGRAD_CLIP_NORM = 1.0\nTTA_VIEWS      = 5     # number of augmented views for TTA\nDEVICE         = 'cuda' if torch.cuda.is_available() else 'cpu'\nUSE_AMP        = DEVICE == 'cuda'\nVERSION        = 'v05'\nCV_FOLD        = 0\n\nprint('\\n' + '='*60)\nprint(f'  {VERSION}: Augmentation + TTA')\nprint(f'  Backbone   : {BACKBONE}')\nprint(f'  ArcFace    : m={ARC_MARGIN}, s={ARC_SCALE}')\nprint(f'  Aug level  : heavy')\nprint(f'  TTA views  : {TTA_VIEWS}')\nprint(f'  Device     : {DEVICE}  |  AMP: {USE_AMP}')\nprint(f'  Epochs     : S1={EPOCHS_S1} + S2={EPOCHS_S2}')\nprint('='*60)\n\n# -- Load data -----------------------------------------------------------------\nprint('\\n[1] 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'])\n\nfor fold, train_idx, val_idx in get_identity_splits(train_df, n_splits=5):\n    if fold == CV_FOLD:\n        break\n\n# -- Pre-cache images ----------------------------------------------------------\nprint('\\n[2] Pre-caching images...')\nall_train_files = train_df['filename'].tolist()\ntest_files_flat = sorted(pd.unique(test_df[['query_image', 'gallery_image']].values.ravel()))\ntrain_cache = build_image_cache(all_train_files, TRAIN_DIR, img_size=IMG_SIZE, verbose=True)\ntest_cache  = build_image_cache(test_files_flat,  TEST_DIR,  img_size=IMG_SIZE, verbose=True)\n\n# Heavy augmentation for training; clean val transforms\ntf_train = get_transforms(mode='train', img_size=IMG_SIZE, augment_level='heavy', cached=True)\ntf_val   = get_transforms(mode='val',   img_size=IMG_SIZE, cached=True)\n\ntrain_ds = JaguarDataset(\n    filenames=train_df['filename'].iloc[train_idx].tolist(),\n    img_dir=TRAIN_DIR, labels=labels[train_idx],\n    transform=tf_train, image_cache=train_cache)\nval_ds = JaguarDataset(\n    filenames=train_df['filename'].iloc[val_idx].tolist(),\n    img_dir=TRAIN_DIR, labels=labels[val_idx],\n    transform=tf_val, image_cache=train_cache)\nfull_ds = JaguarDataset(\n    filenames=all_train_files, img_dir=TRAIN_DIR, labels=labels,\n    transform=tf_val, image_cache=train_cache)\nfull_train_ds = JaguarDataset(\n    filenames=all_train_files, img_dir=TRAIN_DIR, labels=labels,\n    transform=tf_train, image_cache=train_cache)\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)} | Val: {len(val_ds)}')\n\n\n# -- TTA utilities -------------------------------------------------------------\n\ndef get_tta_transforms(img_size: int):\n    \"\"\"\n    Returns a list of deterministic transforms for TTA.\n    All produce (img_size, img_size) tensors from a (img_size+32) cached PIL image.\n\n    Views:\n      0: centre crop (same as val)\n      1: horizontal flip + centre crop\n      2: top-left crop\n      3: top-right crop\n      4: bottom-centre crop\n    \"\"\"\n    norm = [transforms.ToTensor(), transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD)]\n    pad  = img_size + 32  # cached size\n    off  = 32             # = pad - img_size; offset for corner crops\n\n    return [\n        transforms.Compose([transforms.CenterCrop(img_size)] + norm),\n        transforms.Compose([transforms.RandomHorizontalFlip(p=1.0),\n                             transforms.CenterCrop(img_size)] + norm),\n        transforms.Compose([transforms.Lambda(lambda img: img.crop((0, 0, img_size, img_size)))] + norm),\n        transforms.Compose([transforms.Lambda(lambda img: img.crop((off, 0, pad, img_size)))] + norm),\n        transforms.Compose([transforms.Lambda(lambda img: img.crop((off//2, off, off//2+img_size, pad)))] + norm),\n    ]\n\n\ndef extract_embeddings_tta(model, filenames, img_dir, cache, img_size,\n                            batch_size, device):\n    \"\"\"\n    Extract L2-normalised embeddings with TTA.\n\n    For each image, runs TTA_VIEWS augmented views through the model and\n    averages the resulting unit-norm embeddings. The average is re-normalised\n    to stay on the hypersphere.\n\n    Returns: (N, embedding_dim) numpy array\n    \"\"\"\n    tta_tfs = get_tta_transforms(img_size)\n    model.eval()\n\n    all_embs = []   # will be (N, TTA_VIEWS, D)\n\n    for tf in tta_tfs:\n        ds = JaguarDataset(filenames=filenames, img_dir=img_dir,\n                           labels=None, transform=tf, image_cache=cache)\n        loader = DataLoader(ds, batch_size=batch_size, shuffle=False,\n                            num_workers=2, pin_memory=True)\n        view_embs = []\n        with torch.no_grad():\n            for imgs in loader:\n                imgs = imgs.to(device)\n                with autocast(enabled=USE_AMP):\n                    embs = model(imgs)   # already L2-normalised\n                view_embs.append(embs.cpu().numpy())\n        all_embs.append(np.concatenate(view_embs, axis=0))\n\n    # Stack: (TTA_VIEWS, N, D) -> mean -> (N, D) -> re-normalise\n    stacked = np.stack(all_embs, axis=0)   # (V, N, D)\n    avg     = stacked.mean(axis=0)         # (N, D)\n    norms   = np.linalg.norm(avg, axis=1, keepdims=True).clip(min=1e-8)\n    return avg / norms\n\n\n# -- Model ---------------------------------------------------------------------\n\nclass ArcFaceModel(nn.Module):\n    def __init__(self, backbone_name, embedding_dim, pretrained=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):\n        return F.normalize(self.proj(self.backbone(x)), p=2, dim=1)\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\nprint('\\n[3] Building model...')\nmodel   = ArcFaceModel(BACKBONE, EMBEDDING_DIM).to(DEVICE)\narcface = ArcFaceLoss(EMBEDDING_DIM, N_CLASSES, ARC_MARGIN, ARC_SCALE).to(DEVICE)\nscaler  = GradScaler(enabled=USE_AMP)\nprint(f'  Backbone dim: {model.backbone.out_dim} | '\n      f'Total params: {sum(p.numel() for p in model.parameters()):,}')\n\n\n# -- Training ------------------------------------------------------------------\n\ndef run_epoch(model, arcface_loss, loader, optimizer, scheduler,\n              is_train, device, scaler):\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            with autocast(enabled=USE_AMP):\n                loss = arcface_loss(model(imgs), targets)\n            if is_train:\n                scaler.scale(loss).backward()\n                scaler.unscale_(optimizer)\n                torch.nn.utils.clip_grad_norm_(\n                    list(model.parameters()) + list(arcface_loss.parameters()),\n                    max_norm=GRAD_CLIP_NORM)\n                scaler.step(optimizer)\n                scaler.update()\n                scheduler.step()\n            total_loss += loss.item() * len(imgs)\n            n          += len(imgs)\n    return total_loss / n\n\n\nprint('\\n[4] Stage 1: head only...')\nmodel.freeze_backbone()\nhead_params = list(filter(lambda p: p.requires_grad, model.parameters())) + \\\n              list(arcface.parameters())\nopt1 = AdamW(head_params, lr=LR_HEAD, weight_decay=WEIGHT_DECAY)\nsch1 = OneCycleLR(opt1, max_lr=LR_HEAD,\n                  steps_per_epoch=len(train_loader), epochs=EPOCHS_S1)\nfor ep in range(1, EPOCHS_S1 + 1):\n    t0 = time.time()\n    tr = run_epoch(model, arcface, train_loader, opt1, sch1, True,  DEVICE, scaler)\n    va = run_epoch(model, arcface, val_loader,   opt1, sch1, False, DEVICE, scaler)\n    print(f'  S1 ep {ep}/{EPOCHS_S1} | loss {tr:.4f}/{va:.4f} | {time.time()-t0:.0f}s')\n\nprint('\\n[5] Stage 2: full fine-tune...')\nmodel.unfreeze_backbone()\nopt2 = AdamW([\n    {'params': model.backbone.parameters(), 'lr': LR_BACKBONE},\n    {'params': model.proj.parameters(),     'lr': LR_HEAD},\n    {'params': arcface.parameters(),        'lr': LR_HEAD},\n], weight_decay=WEIGHT_DECAY)\nsch2 = OneCycleLR(opt2, max_lr=[LR_BACKBONE, LR_HEAD, LR_HEAD],\n                  steps_per_epoch=len(train_loader), epochs=EPOCHS_S2)\n\nbest_val_map = 0.0\nbest_state   = None\n\nfor ep in range(1, EPOCHS_S2 + 1):\n    t0 = time.time()\n    tr = run_epoch(model, arcface, train_loader, opt2, sch2, True,  DEVICE, scaler)\n    va = run_epoch(model, arcface, val_loader,   opt2, sch2, False, DEVICE, scaler)\n\n    if ep % 5 == 0 or ep == EPOCHS_S2:\n        # Val mAP without TTA (fast)\n        val_embs    = []\n        model.eval()\n        with torch.no_grad():\n            for imgs, _ in val_loader:\n                with autocast(enabled=USE_AMP):\n                    val_embs.append(model(imgs.to(DEVICE)).cpu().numpy())\n        val_embs    = np.concatenate(val_embs, axis=0)\n        val_results = compute_map(val_embs, labels[val_idx])\n        vm = val_results['map']\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 = ' **'\n        print(f'  S2 ep {ep}/{EPOCHS_S2} | loss {tr:.4f}/{va:.4f} | '\n              f'mAP {vm:.4f}{marker} | {time.time()-t0:.0f}s')\n    else:\n        print(f'  S2 ep {ep}/{EPOCHS_S2} | loss {tr:.4f}/{va:.4f} | {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: val fold without TTA and with TTA -------------------------------\nprint('\\n[6] Evaluating val fold (no TTA vs TTA)...')\n\n# Without TTA\nval_embs_plain = []\nmodel.eval()\nwith torch.no_grad():\n    for imgs, _ in val_loader:\n        with autocast(enabled=USE_AMP):\n            val_embs_plain.append(model(imgs.to(DEVICE)).cpu().numpy())\nval_embs_plain = np.concatenate(val_embs_plain, axis=0)\nres_plain = compute_map(val_embs_plain, labels[val_idx])\nprint(f'  Val mAP (no TTA) : {res_plain[\"map\"]:.4f}')\n\n# With TTA\nval_filenames = train_df['filename'].iloc[val_idx].tolist()\nval_embs_tta  = extract_embeddings_tta(\n    model, val_filenames, TRAIN_DIR, train_cache, IMG_SIZE, BATCH_SIZE, DEVICE)\nres_tta = compute_map(val_embs_tta, labels[val_idx])\nprint(f'  Val mAP (TTA x{TTA_VIEWS}): {res_tta[\"map\"]:.4f}  '\n      f'(delta: {res_tta[\"map\"] - res_plain[\"map\"]:+.4f})')\n\nprint_results(res_tta, idx_to_label, title=f'{VERSION} - Fold {CV_FOLD} Val (TTA)')\n\n# Full train mAP for reference\nfull_embs    = []\nfull_loader_eval = DataLoader(full_ds, batch_size=BATCH_SIZE, shuffle=False,\n                               num_workers=2, pin_memory=True)\nmodel.eval()\nwith torch.no_grad():\n    for imgs, _ in full_loader_eval:\n        with autocast(enabled=USE_AMP):\n            full_embs.append(model(imgs.to(DEVICE)).cpu().numpy())\nfull_embs    = np.concatenate(full_embs, axis=0)\nfull_results = compute_map(full_embs, labels)\nprint(f'\\n  Full train mAP (no TTA): {full_results[\"map\"]:.4f}')\nprint(f'  v04 ConvNeXt was 0.8382 — delta: {full_results[\"map\"]-0.8382:+.4f}')\n\nsave_benchmark(\n    OUT_DIR / 'benchmarks_v05.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':   'heavy',\n        'cv_map_val':     round(res_tta['map'], 5),\n        'cv_map_val_noTTA': round(res_plain['map'], 5),\n        'cv_map_full':    round(full_results['map'], 5),\n        'best_identity':  idx_to_label[max(res_tta['per_identity_ap'],\n                                           key=res_tta['per_identity_ap'].get)],\n        'worst_identity': idx_to_label[min(res_tta['per_identity_ap'],\n                                           key=res_tta['per_identity_ap'].get)],\n        'epochs_s1':      EPOCHS_S1,\n        'epochs_s2':      EPOCHS_S2,\n        'notes':          f'ConvNeXt-Small + heavy aug + TTA x{TTA_VIEWS}',\n    }\n)\n\n\n# -- Visualisations ------------------------------------------------------------\nprint('\\n[7] Generating visualisations...')\n\n# Per-identity AP: TTA vs no-TTA vs v04\nv05_per_id = {idx_to_label[k]: v for k, v in res_tta['per_identity_ap'].items()}\nv04_per_id_ref = {idx_to_label[k]: v for k, v in res_plain['per_identity_ap'].items()}\nids_sorted = sorted(v05_per_id.keys(), key=lambda x: -v05_per_id[x])\nx = np.arange(len(ids_sorted))\nw = 0.38\n\nfig, ax = plt.subplots(figsize=(16, 5))\nax.bar(x - w/2, [v04_per_id_ref.get(i, 0) for i in ids_sorted], w,\n       color='lightsteelblue', label='v05 no TTA', edgecolor='white')\nax.bar(x + w/2, [v05_per_id[i] for i in ids_sorted], w,\n       color='steelblue', label=f'v05 TTA x{TTA_VIEWS}', edgecolor='white')\nax.axhline(res_tta['map'], color='navy', linestyle='--', linewidth=1,\n           label=f'v05 TTA mAP={res_tta[\"map\"]:.4f}')\nax.axhline(res_plain['map'], color='gray', linestyle=':', linewidth=1,\n           label=f'v05 no-TTA mAP={res_plain[\"map\"]:.4f}')\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('v05: Per-Identity AP — TTA vs No TTA', 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 -> v05_per_identity_ap.png')\n\n# t-SNE\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'v05: t-SNE — ConvNeXt-Small + Heavy Aug', 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 -> v05_tsne.png')\n\n\n# -- Full retrain + TTA submission ---------------------------------------------\nprint('\\n[8] Retraining on ALL data + TTA submission...')\n\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, ARC_MARGIN, ARC_SCALE).to(DEVICE)\nscaler_final  = GradScaler(enabled=USE_AMP)\n\n# Stage 1\nmodel_final.freeze_backbone()\nhead_f = list(filter(lambda p: p.requires_grad, model_final.parameters())) + \\\n         list(arcface_final.parameters())\nopt_f1 = AdamW(head_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,\n                     opt_f1, sch_f1, True, DEVICE, scaler_final)\n    print(f'  Full S1 ep {ep}/{EPOCHS_S1} | loss {loss:.4f}')\n\n# Stage 2\nmodel_final.unfreeze_backbone()\nopt_f2 = AdamW([\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], weight_decay=WEIGHT_DECAY)\nsch_f2 = OneCycleLR(opt_f2, max_lr=[LR_BACKBONE, LR_HEAD, LR_HEAD],\n                    steps_per_epoch=len(full_loader), epochs=EPOCHS_S2)\nfor ep in range(1, EPOCHS_S2 + 1):\n    loss = run_epoch(model_final, arcface_final, full_loader,\n                     opt_f2, sch_f2, True, DEVICE, scaler_final)\n    if ep % 5 == 0:\n        print(f'  Full S2 ep {ep}/{EPOCHS_S2} | loss {loss:.4f}')\n\n# TTA submission\nprint(f'  Extracting test embeddings with TTA x{TTA_VIEWS}...')\nmodel_final.eval()\ntest_embs_tta = extract_embeddings_tta(\n    model_final, test_files_flat, TEST_DIR, test_cache, IMG_SIZE, BATCH_SIZE, DEVICE)\nfname_to_idx  = {f: i for i, f in enumerate(test_files_flat)}\nmake_submission(\n    test_df=test_df, test_embeddings=test_embs_tta,\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')\n\nprint(f'\\n{\"=\"*60}')\nprint(f'  {VERSION} COMPLETE')\nprint(f'  Val mAP no-TTA: {res_plain[\"map\"]:.4f}')\nprint(f'  Val mAP TTA x{TTA_VIEWS}: {res_tta[\"map\"]:.4f}')\nprint(f'  Full train mAP : {full_results[\"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":"ab9cce0f-e5eb-46e7-acf0-45ca86bb9183","_cell_guid":"8956d74e-2970-409b-abfa-fbc294e1e42e","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null}]}