{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","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}],"dockerImageVersionId":31288,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"md_header","cell_type":"markdown","source":"# 🐆 Jaguar Re-ID: The Winning Blueprint\n\n> **Competition:** Jaguar Re-Identification Challenge (Pantanal)  \n> **Task:** Pairwise similarity prediction for 371 test images (137,270 pairs)  \n> **Metric:** Identity-Balanced mAP  \n\n---\n\n## Pipeline Overview\n\n| Component | Details |\n|-----------|---------|\n| Backbone | EfficientNetV2 (auto-selected by hardware) |\n| Loss | ArcFace (s=30, m=0.5) |\n| Pooling | GeM (Generalized Mean Pooling) |\n| Augmentation | Flip + Brightness + ShiftScaleRotate + CoarseDropout |\n| TTA | 3 variants (original + flip + crop) |\n| Optimizer | AdamW + CosineAnnealingLR |\n| Hardware | Auto-detects GPU/CPU and adjusts settings |\n\n---\n","metadata":{}},{"id":"md_install","cell_type":"markdown","source":"## ⚙️ Step 1 — Install Dependencies","metadata":{}},{"id":"install","cell_type":"code","source":"import subprocess, sys\n\ndef install(pkg):\n    subprocess.run([sys.executable, '-m', 'pip', 'install', pkg, '-q'],\n                   capture_output=True)\n\ninstall('timm')\ninstall('albumentations')\ninstall('umap-learn')\ninstall('seaborn')\ntry:\n    subprocess.run([sys.executable, '-m', 'pip', 'install', 'faiss-gpu', '-q'],\n                   capture_output=True, check=True)\n    print(\"faiss-gpu installed\")\nexcept:\n    install('faiss-cpu')\n    print(\"faiss-cpu installed (fallback)\")\n\nprint(\"All dependencies ready!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-14T05:43:15.68928Z","iopub.execute_input":"2026-03-14T05:43:15.689823Z","iopub.status.idle":"2026-03-14T05:43:32.457744Z","shell.execute_reply.started":"2026-03-14T05:43:15.689784Z","shell.execute_reply":"2026-03-14T05:43:32.456747Z"}},"outputs":[],"execution_count":null},{"id":"md_imports","cell_type":"markdown","source":"## 📦 Step 2 — Imports & Device Detection","metadata":{}},{"id":"imports","cell_type":"code","source":"import os, gc, math, random, time, warnings\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader, WeightedRandomSampler\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom torch.cuda.amp import autocast, GradScaler\nfrom pathlib import Path\nfrom PIL import Image\nfrom tqdm.auto import tqdm\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm\nimport matplotlib.pyplot as plt\nimport matplotlib.gridspec as gridspec\nimport seaborn as sns\nwarnings.filterwarnings('ignore')\n\n# ── Auto device detection ──────────────────────────────\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nAMP    = torch.cuda.is_available()\nN_GPU  = torch.cuda.device_count()\n\ndef set_seed(seed=42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if AMP:\n        torch.cuda.manual_seed_all(seed)\n\nset_seed(42)\n\nprint(f\"Device  : {DEVICE}\")\nprint(f\"GPUs    : {N_GPU}\")\nprint(f\"AMP fp16: {AMP}\")\nif AMP:\n    for i in range(N_GPU):\n        props = torch.cuda.get_device_properties(i)\n        print(f\"  GPU {i}: {props.name} ({props.total_memory//1024**2} MB)\")\nelse:\n    import multiprocessing\n    print(f\"CPU cores: {multiprocessing.cpu_count()}\")\n    print(\"NOTE: Running on CPU — reduced settings will be applied automatically\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-14T05:43:32.458857Z","iopub.execute_input":"2026-03-14T05:43:32.459175Z","iopub.status.idle":"2026-03-14T05:43:40.580759Z","shell.execute_reply.started":"2026-03-14T05:43:32.459141Z","shell.execute_reply":"2026-03-14T05:43:40.579901Z"}},"outputs":[],"execution_count":null},{"id":"md_cfg","cell_type":"markdown","source":"## 🧠 Step 3 — Smart Configuration (Auto-adapts to Hardware)","metadata":{}},{"id":"cfg","cell_type":"code","source":"# ── Auto-detect data path ─────────────────────────────\ndef find_data():\n    candidates = [\n        Path('/kaggle/input/competitions/jaguar-re-id'),\n        Path('/kaggle/input/jaguar-re-id'),\n        Path('/kaggle/input/jaguar-re-identification'),\n        Path('/kaggle/input/jaguar-re-identification-challenge'),\n    ]\n    for p in candidates:\n        if (p / 'train.csv').exists():\n            return p\n    for root, dirs, files in os.walk('/kaggle/input'):\n        if 'train.csv' in files:\n            return Path(root)\n    return Path('.')\n\nDATA_DIR = find_data()\n\n# ── Hardware-aware config ──────────────────────────────\nif AMP and N_GPU >= 2:\n    BATCH_SIZE, EPOCHS, NUM_WORKERS = 64, 50, 4\n    IMG_SIZE, BACKBONE = 224, 'tf_efficientnetv2_m'\n    print(\"Mode: HIGH PERFORMANCE (Multi-GPU)\")\nelif AMP and N_GPU == 1:\n    BATCH_SIZE, EPOCHS, NUM_WORKERS = 32, 40, 2\n    IMG_SIZE, BACKBONE = 224, 'tf_efficientnetv2_s'\n    print(\"Mode: STANDARD GPU\")\nelse:\n    BATCH_SIZE, EPOCHS, NUM_WORKERS = 8, 5, 0\n    IMG_SIZE, BACKBONE = 128, 'efficientnet_b0'\n    print(\"Mode: CPU FALLBACK (demo mode)\")\n\nEMBED_DIM  = 512\nNUM_CLS    = 31\nLR         = 3e-4\nOUT_DIR    = Path('/kaggle/working')\nTRAIN_DIR  = DATA_DIR / 'train' / 'train'\nTEST_DIR   = DATA_DIR / 'test'  / 'test'\nTRAIN_CSV  = DATA_DIR / 'train.csv'\nTEST_CSV   = DATA_DIR / 'test.csv'\nSAMPLE_SUB = DATA_DIR / 'sample_submission.csv'\n\nif not TRAIN_DIR.exists(): TRAIN_DIR = DATA_DIR / 'train'\nif not TEST_DIR.exists():  TEST_DIR  = DATA_DIR / 'test'\nOUT_DIR.mkdir(exist_ok=True)\n\nprint(f\"Data dir   : {DATA_DIR}\")\nprint(f\"Backbone   : {BACKBONE}\")\nprint(f\"Batch size : {BATCH_SIZE}\")\nprint(f\"Epochs     : {EPOCHS}\")\nprint(f\"Image size : {IMG_SIZE}\")\nprint(f\"Train dir  : {TRAIN_DIR} ({'OK' if TRAIN_DIR.exists() else 'MISSING'})\")\nprint(f\"Test dir   : {TEST_DIR}  ({'OK' if TEST_DIR.exists() else 'MISSING'})\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-14T05:43:40.58266Z","iopub.execute_input":"2026-03-14T05:43:40.583082Z","iopub.status.idle":"2026-03-14T05:43:40.600863Z","shell.execute_reply.started":"2026-03-14T05:43:40.583056Z","shell.execute_reply":"2026-03-14T05:43:40.600023Z"}},"outputs":[],"execution_count":null},{"id":"md_eda","cell_type":"markdown","source":"## 🔍 Step 4 — Exploratory Data Analysis","metadata":{}},{"id":"eda","cell_type":"code","source":"train_df = pd.read_csv(TRAIN_CSV)\ntest_df  = pd.read_csv(TEST_CSV)\nsub_df   = pd.read_csv(SAMPLE_SUB)\n\nlabels   = sorted(train_df['ground_truth'].unique())\nl2i      = {l: i for i, l in enumerate(labels)}\ntrain_df['label'] = train_df['ground_truth'].map(l2i)\n\nprint(f\"Train images : {len(train_df)}\")\nprint(f\"Test images  : {len(test_df)}\")\nprint(f\"Identities   : {train_df['ground_truth'].nunique()}\")\nprint(f\"Test pairs   : {len(sub_df):,}\")\nprint(f\"Avg per ID   : {len(train_df)/train_df['ground_truth'].nunique():.1f}\")\n\nplt.style.use('dark_background')\nfig, axes = plt.subplots(1, 3, figsize=(20, 5), facecolor='#0a0e17')\n\ncounts = train_df['ground_truth'].value_counts()\n\nax = axes[0]\nax.set_facecolor('#0f1520')\nax.bar(range(len(counts)), counts.values,\n       color=plt.cm.YlOrBr(np.linspace(0.3, 0.9, len(counts))), edgecolor='#0a0e17')\nax.axhline(counts.mean(), color='#ff6b35', linestyle='--', linewidth=2,\n           label=f'Mean={counts.mean():.1f}')\nax.set_title('Images per Jaguar Identity', color='#d4a853', fontsize=12, fontweight='bold')\nax.set_xlabel('Jaguar ID', color='#8b9ab0')\nax.set_ylabel('Count', color='#8b9ab0')\nax.tick_params(colors='#8b9ab0')\nfor s in ax.spines.values(): s.set_color('#1a2030')\nax.legend(facecolor='#1a2030', labelcolor='white')\n\nax2 = axes[1]\nax2.set_facecolor('#0f1520')\nax2.hist(counts.values, bins=12, color='#4ecdc4', edgecolor='#0a0e17', alpha=0.85)\nax2.set_title('Distribution of Images per Identity', color='#d4a853', fontsize=12, fontweight='bold')\nax2.set_xlabel('Images per jaguar', color='#8b9ab0')\nax2.set_ylabel('Frequency', color='#8b9ab0')\nax2.tick_params(colors='#8b9ab0')\nfor s in ax2.spines.values(): s.set_color('#1a2030')\n\nax3 = axes[2]\nax3.set_facecolor('#0f1520')\nax3.axis('off')\nstats = [\n    ['Train images', f'{len(train_df):,}'],\n    ['Test images', f'{len(test_df):,}'],\n    ['Unique jaguars', f'{train_df[\"ground_truth\"].nunique()}'],\n    ['Test pairs', f'{len(sub_df):,}'],\n    ['Min per ID', f'{counts.min()}'],\n    ['Max per ID', f'{counts.max()}'],\n    ['Mean per ID', f'{counts.mean():.1f}'],\n]\nt = ax3.table(cellText=stats, colLabels=['Metric', 'Value'],\n              cellLoc='center', loc='center',\n              cellColours=[['#1a2030','#1a2030']]*len(stats),\n              colColours=['#d4a853','#d4a853'])\nt.auto_set_font_size(False)\nt.set_fontsize(11)\nt.scale(1.3, 1.9)\nfor (r,c), cell in t.get_celld().items():\n    cell.set_text_props(color='white' if r>0 else '#0a0e17', fontweight='bold' if r==0 else 'normal')\n    cell.set_edgecolor('#1a2030')\nax3.set_title('Dataset Statistics', color='#d4a853', fontsize=12, fontweight='bold', pad=10)\n\nplt.suptitle('Jaguar Re-ID — Dataset Overview', color='#d4a853', fontsize=15, fontweight='bold')\nplt.tight_layout()\nplt.savefig(OUT_DIR/'eda.png', dpi=130, bbox_inches='tight', facecolor='#0a0e17')\nplt.show()\nprint(\"EDA complete!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-14T05:43:40.601976Z","iopub.execute_input":"2026-03-14T05:43:40.602504Z","iopub.status.idle":"2026-03-14T05:43:41.766614Z","shell.execute_reply.started":"2026-03-14T05:43:40.602472Z","shell.execute_reply":"2026-03-14T05:43:41.765856Z"}},"outputs":[],"execution_count":null},{"id":"md_transforms","cell_type":"markdown","source":"## 🔄 Step 5 — Augmentation Pipelines","metadata":{}},{"id":"transforms","cell_type":"code","source":"train_tfm = A.Compose([\n    A.Resize(IMG_SIZE, IMG_SIZE),\n    A.HorizontalFlip(p=0.5),\n    A.RandomBrightnessContrast(0.3, 0.3, p=0.6),\n    A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=15, p=0.5),\n    A.CoarseDropout(max_holes=4, max_height=32, max_width=32, p=0.3),\n    A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n    ToTensorV2(),\n])\n\nval_tfm = A.Compose([\n    A.Resize(IMG_SIZE, IMG_SIZE),\n    A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n    ToTensorV2(),\n])\n\ntta_tfms = [\n    A.Compose([A.Resize(IMG_SIZE, IMG_SIZE),\n               A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n               ToTensorV2()]),\n    A.Compose([A.Resize(IMG_SIZE, IMG_SIZE), A.HorizontalFlip(p=1.0),\n               A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n               ToTensorV2()]),\n    A.Compose([A.Resize(int(IMG_SIZE*1.1), int(IMG_SIZE*1.1)),\n               A.CenterCrop(IMG_SIZE, IMG_SIZE),\n               A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n               ToTensorV2()]),\n]\nprint(f\"Train augmentations : {len(train_tfm.transforms)}\")\nprint(f\"TTA variants        : {len(tta_tfms)}\")\nprint(f\"Image size          : {IMG_SIZE}x{IMG_SIZE}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-14T05:43:41.767571Z","iopub.execute_input":"2026-03-14T05:43:41.767902Z","iopub.status.idle":"2026-03-14T05:43:41.78439Z","shell.execute_reply.started":"2026-03-14T05:43:41.767877Z","shell.execute_reply":"2026-03-14T05:43:41.783803Z"}},"outputs":[],"execution_count":null},{"id":"md_datasets","cell_type":"markdown","source":"## 🗂️ Step 6 — Dataset Classes","metadata":{}},{"id":"datasets","cell_type":"code","source":"class TrainDS(Dataset):\n    def __init__(self, df, img_dir, tfm=None):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = Path(img_dir)\n        self.tfm = tfm\n    def __len__(self):\n        return len(self.df)\n    def __getitem__(self, i):\n        row = self.df.iloc[i]\n        img = np.array(Image.open(self.img_dir/row['filename']).convert('RGB'))\n        if self.tfm:\n            img = self.tfm(image=img)['image']\n        return img, torch.tensor(row['label'], dtype=torch.long)\n\n\nclass TestDS(Dataset):\n    def __init__(self, files, img_dir, tfm=None):\n        self.files   = list(files)\n        self.img_dir = Path(img_dir)\n        self.tfm     = tfm\n    def __len__(self):\n        return len(self.files)\n    def __getitem__(self, i):\n        f   = self.files[i]\n        img = np.array(Image.open(self.img_dir/f).convert('RGB'))\n        if self.tfm:\n            img = self.tfm(image=img)['image']\n        return img, f\n\n\ndef get_sampler(df):\n    counts  = df['label'].value_counts().sort_index()\n    weights = 1.0 / counts[df['label'].values].values\n    return WeightedRandomSampler(torch.FloatTensor(weights), len(df), replacement=True)\n\n\nprint(\"Dataset classes ready!\")\nprint(f\"  TrainDS — with balanced sampler\")\nprint(f\"  TestDS  — for inference\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-14T05:43:41.785465Z","iopub.execute_input":"2026-03-14T05:43:41.785844Z","iopub.status.idle":"2026-03-14T05:43:41.796812Z","shell.execute_reply.started":"2026-03-14T05:43:41.785813Z","shell.execute_reply":"2026-03-14T05:43:41.796065Z"}},"outputs":[],"execution_count":null},{"id":"md_model","cell_type":"markdown","source":"## 🧠 Step 7 — Model Architecture (ArcFace + GeM)","metadata":{}},{"id":"model","cell_type":"code","source":"class ArcFace(nn.Module):\n    def __init__(self, dim, n_cls, s=30.0, m=0.5):\n        super().__init__()\n        self.s  = s\n        self.m  = m\n        self.w  = nn.Parameter(torch.FloatTensor(n_cls, dim))\n        nn.init.xavier_uniform_(self.w)\n        self.cm = math.cos(m)\n        self.sm = math.sin(m)\n        self.th = math.cos(math.pi - m)\n        self.mm = math.sin(math.pi - m) * m\n\n    def forward(self, f, y):\n        c  = F.linear(F.normalize(f), F.normalize(self.w))\n        s  = torch.sqrt((1 - c.pow(2)).clamp(1e-9))\n        p  = c * self.cm - s * self.sm\n        p  = torch.where(c > self.th, p, c - self.mm)\n        oh = torch.zeros_like(c).scatter_(1, y.view(-1,1), 1)\n        return F.cross_entropy((oh * p + (1-oh) * c) * self.s, y)\n\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\n    def forward(self, x):\n        return F.adaptive_avg_pool2d(\n            x.clamp(min=self.eps).pow(self.p), (1,1)\n        ).pow(1.0 / self.p)\n\n\nclass JaguarModel(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.backbone = timm.create_model(\n            BACKBONE, pretrained=True, num_classes=0, global_pool=''\n        )\n        in_feat   = self.backbone.num_features\n        self.pool = GeM()\n        self.neck = nn.Sequential(\n            nn.Dropout(0.3),\n            nn.Linear(in_feat, EMBED_DIM),\n            nn.BatchNorm1d(EMBED_DIM),\n        )\n\n    def forward(self, x):\n        f = self.backbone(x)\n        if f.dim() == 4:\n            f = self.pool(f).flatten(1)\n        elif f.dim() == 3:\n            f = f[:, 0]\n        else:\n            f = f.flatten(1)\n        return F.normalize(self.neck(f), p=2, dim=1)\n\n\ngc.collect()\nif AMP:\n    torch.cuda.empty_cache()\n\nmodel   = JaguarModel().to(DEVICE)\narcface = ArcFace(EMBED_DIM, NUM_CLS).to(DEVICE)\n\nif N_GPU > 1:\n    model = nn.DataParallel(model)\n    print(f\"DataParallel enabled on {N_GPU} GPUs\")\n\ntotal_p = sum(p.numel() for p in model.parameters()) / 1e6\nprint(f\"Backbone    : {BACKBONE}\")\nprint(f\"Parameters  : {total_p:.1f}M\")\nprint(f\"Embed dim   : {EMBED_DIM}\")\nprint(f\"Num classes : {NUM_CLS}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-14T05:43:41.797717Z","iopub.execute_input":"2026-03-14T05:43:41.798035Z","iopub.status.idle":"2026-03-14T05:43:44.868463Z","shell.execute_reply.started":"2026-03-14T05:43:41.79801Z","shell.execute_reply":"2026-03-14T05:43:44.867728Z"}},"outputs":[],"execution_count":null},{"id":"md_training","cell_type":"markdown","source":"## 🏋️ Step 8 — Training Loop","metadata":{}},{"id":"training","cell_type":"code","source":"ds      = TrainDS(train_df, TRAIN_DIR, train_tfm)\nsampler = get_sampler(train_df)\nloader  = DataLoader(ds, batch_size=BATCH_SIZE, sampler=sampler,\n                     num_workers=NUM_WORKERS, pin_memory=AMP, drop_last=True)\n\noptimizer = AdamW(\n    list(model.parameters()) + list(arcface.parameters()),\n    lr=LR, weight_decay=1e-4\n)\nscheduler = CosineAnnealingLR(optimizer, T_max=EPOCHS, eta_min=1e-6)\nscaler    = GradScaler(enabled=AMP)\n\nbest_loss = float('inf')\nckpt_path = OUT_DIR / 'best_model.pth'\nhistory   = []\n\nprint(f\"Training started!\")\nprint(f\"  Epochs     : {EPOCHS}\")\nprint(f\"  Batch size : {BATCH_SIZE}\")\nprint(f\"  Steps/epoch: {len(loader)}\")\nprint(f\"  AMP fp16   : {AMP}\")\nprint(\"-\" * 45)\nstart = time.time()\n\nfor epoch in range(EPOCHS):\n    model.train()\n    arcface.train()\n    total = 0\n    pbar  = tqdm(loader, desc=f'Epoch {epoch+1:02d}/{EPOCHS}')\n\n    for imgs, labels_ in pbar:\n        imgs    = imgs.to(DEVICE, non_blocking=True)\n        labels_ = labels_.to(DEVICE, non_blocking=True)\n        with autocast(enabled=AMP):\n            loss = arcface(model(imgs), labels_)\n        optimizer.zero_grad()\n        scaler.scale(loss).backward()\n        scaler.unscale_(optimizer)\n        nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n        scaler.step(optimizer)\n        scaler.update()\n        total += loss.item()\n        pbar.set_postfix(loss=f'{loss.item():.4f}')\n\n    scheduler.step()\n    avg     = total / len(loader)\n    elapsed = (time.time() - start) / 60\n    history.append(avg)\n    print(f\"Epoch {epoch+1:02d}/{EPOCHS} | Loss: {avg:.4f} | Time: {elapsed:.1f}m\")\n\n    if avg < best_loss:\n        best_loss = avg\n        m_state   = model.module.state_dict() if N_GPU > 1 else model.state_dict()\n        torch.save({'model': m_state, 'backbone': BACKBONE,\n                    'embed_dim': EMBED_DIM, 'img_size': IMG_SIZE}, ckpt_path)\n        print(f\"  Saved best checkpoint (loss={best_loss:.4f})\")\n\ntotal_time = (time.time() - start) / 60\nprint(f\"Training complete! Best loss: {best_loss:.4f} | Total: {total_time:.1f}m\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-14T05:43:44.869516Z","iopub.execute_input":"2026-03-14T05:43:44.870078Z"}},"outputs":[],"execution_count":null},{"id":"md_inference","cell_type":"markdown","source":"## 🔬 Step 9 — TTA Inference","metadata":{}},{"id":"tta_inference","cell_type":"code","source":"def extract_embeddings(files, img_dir, use_tta=True):\n    m    = model.module if N_GPU > 1 else model\n    m.eval()\n    tfms = tta_tfms if (use_tta and AMP) else [val_tfm]\n    all_embs = []\n\n    with torch.no_grad():\n        for tfm in tfms:\n            ds_    = TestDS(files, img_dir, tfm)\n            ldr_   = DataLoader(ds_, batch_size=BATCH_SIZE*2,\n                                num_workers=NUM_WORKERS, pin_memory=AMP)\n            embs   = []\n            for imgs, _ in tqdm(ldr_, desc='Inferring', leave=False):\n                imgs = imgs.to(DEVICE, non_blocking=True)\n                with autocast(enabled=AMP):\n                    e = m(imgs)\n                embs.append(e.cpu())\n            all_embs.append(torch.cat(embs))\n\n    return torch.stack(all_embs).mean(0)\n\n\nall_files = sorted(set(\n    test_df['query_image'].tolist() + test_df['gallery_image'].tolist()\n))\nprint(f\"Total test images : {len(all_files)}\")\nprint(f\"TTA enabled       : {AMP}\")\n\ntest_embs = extract_embeddings(all_files, TEST_DIR, use_tta=AMP)\nfile2idx  = {f: i for i, f in enumerate(all_files)}\nprint(f\"Embeddings shape  : {test_embs.shape}\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"md_submission","cell_type":"markdown","source":"## 📤 Step 10 — Generate Submission","metadata":{}},{"id":"4bf5e643-93af-45ab-a85e-120f9bbb8fc9","cell_type":"code","source":"import pandas as pd\nfrom pathlib import Path\n\nDATA_DIR = Path('/kaggle/input/competitions/jaguar-re-id')\nsub_df   = pd.read_csv(DATA_DIR / 'sample_submission.csv')\nprint(\"Columns:\", sub_df.columns.tolist())\nprint(sub_df.head(5))\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"1eb5ee80-1648-4f98-9545-cb6c7bb4858e","cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom pathlib import Path\nfrom PIL import Image\nfrom tqdm.auto import tqdm\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm, math\n\nDEVICE   = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nAMP      = torch.cuda.is_available()\nDATA_DIR = Path('/kaggle/input/competitions/jaguar-re-id')\nTEST_DIR = DATA_DIR / 'test' / 'test'\nOUT_DIR  = Path('/kaggle/working')\n\ntest_df = pd.read_csv(DATA_DIR / 'test.csv')\nsub_df  = pd.read_csv(DATA_DIR / 'sample_submission.csv')\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.adaptive_avg_pool2d(x.clamp(min=self.eps).pow(self.p),(1,1)).pow(1.0/self.p)\n\nclass JaguarModel(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.backbone = timm.create_model('tf_efficientnetv2_s', pretrained=False, num_classes=0, global_pool='')\n        self.pool = GeM()\n        self.neck = nn.Sequential(nn.Dropout(0.3), nn.Linear(self.backbone.num_features, 512), nn.BatchNorm1d(512))\n    def forward(self, x):\n        f = self.backbone(x)\n        if f.dim() == 4: f = self.pool(f).flatten(1)\n        elif f.dim() == 3: f = f[:, 0]\n        return F.normalize(self.neck(f), p=2, dim=1)\n\nmodel = JaguarModel().to(DEVICE)\nckpt  = torch.load(OUT_DIR/'best_model.pth', map_location=DEVICE)\nmodel.load_state_dict(ckpt['model'])\nmodel.eval()\nprint(f\"Model loaded! Device={DEVICE}\")\n\nval_tfm = A.Compose([\n    A.Resize(224, 224),\n    A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n    ToTensorV2(),\n])\n\nclass TestDS(Dataset):\n    def __init__(self, files, img_dir):\n        self.files = list(files)\n        self.img_dir = Path(img_dir)\n    def __len__(self): return len(self.files)\n    def __getitem__(self, i):\n        f = self.files[i]\n        img = np.array(Image.open(self.img_dir/f).convert('RGB'))\n        return val_tfm(image=img)['image'], f\n\nall_files = sorted(set(test_df['query_image'].tolist() + test_df['gallery_image'].tolist()))\nprint(f\"Images: {len(all_files)}\")\n\nloader = DataLoader(TestDS(all_files, TEST_DIR), batch_size=64, num_workers=2, pin_memory=AMP)\nembs   = []\nwith torch.no_grad():\n    for imgs, _ in tqdm(loader, desc='Extracting'):\n        with torch.cuda.amp.autocast(enabled=AMP):\n            embs.append(model(imgs.to(DEVICE)).cpu())\n\nE        = torch.cat(embs).numpy()\nfile2idx = {f: i for i, f in enumerate(all_files)}\nprint(f\"Embeddings: {E.shape}\")\n\nscores = []\nfor _, row in tqdm(test_df.iterrows(), total=len(test_df), desc='Scoring'):\n    qi = file2idx[row['query_image']]\n    gi = file2idx[row['gallery_image']]\n    scores.append(float(np.dot(E[qi], E[gi])))\n\nsub_df['similarity'] = scores\nsub_df.to_csv(OUT_DIR/'submission.csv', index=False)\nprint(f\"Done! shape={sub_df.shape}\")\nprint(f\"Range: [{min(scores):.4f}, {max(scores):.4f}]\")\nprint(sub_df.head())\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"submission","cell_type":"code","source":"print(\"Generating submission...\")\nE_np    = test_embs.numpy()\nresults = []\n\nfor _, row in tqdm(sub_df.iterrows(), total=len(sub_df), desc='Scoring pairs'):\n    qi  = file2idx[row['query_image']]\n    gi  = file2idx[row['gallery_image']]\n    sim = float(np.dot(E_np[qi], E_np[gi]))\n    results.append(sim)\n\nsub_df['score'] = results\nsub_df.to_csv(OUT_DIR/'submission.csv', index=False)\n\nprint(f\"Submission saved!\")\nprint(f\"  Pairs   : {len(sub_df):,}\")\nprint(f\"  Min sim : {sub_df['score'].min():.4f}\")\nprint(f\"  Max sim : {sub_df['score'].max():.4f}\")\nprint(f\"  Mean sim: {sub_df['score'].mean():.4f}\")\nprint(f\"  Std sim : {sub_df['score'].std():.4f}\")\nprint(sub_df.head(3))\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"md_viz","cell_type":"markdown","source":"## 📊 Step 11 — Results Visualization","metadata":{}},{"id":"viz","cell_type":"code","source":"plt.style.use('dark_background')\nfig, axes = plt.subplots(1, 2, figsize=(16, 6), facecolor='#0a0e17')\n\nax = axes[0]\nax.set_facecolor('#0f1520')\nax.plot(range(1, len(history)+1), history,\n        color='#d4a853', linewidth=2.5, marker='o', markersize=4,\n        label=f'Loss (best={min(history):.4f})')\nax.axhline(min(history), color='#ff6b35', linestyle=':', linewidth=1.5, alpha=0.7)\nax.fill_between(range(1, len(history)+1), history, alpha=0.15, color='#d4a853')\nax.set_title('Training Loss Curve', color='#d4a853', fontsize=13, fontweight='bold')\nax.set_xlabel('Epoch', color='#8b9ab0')\nax.set_ylabel('ArcFace Loss', color='#8b9ab0')\nax.tick_params(colors='#8b9ab0')\nfor s in ax.spines.values(): s.set_color('#1a2030')\nax.legend(facecolor='#1a2030', labelcolor='white')\n\nax2 = axes[1]\nax2.set_facecolor('#0f1520')\nax2.hist(sub_df['score'], bins=60, color='#4ecdc4', edgecolor='#0a0e17', alpha=0.85)\nax2.axvline(sub_df['score'].mean(), color='#ff6b35', linestyle='--', linewidth=2,\n            label=f\"Mean={sub_df['score'].mean():.3f}\")\nax2.axvline(0.5, color='#d4a853', linestyle='--', linewidth=2, label='Threshold=0.5')\nax2.set_title('Similarity Score Distribution', color='#d4a853', fontsize=13, fontweight='bold')\nax2.set_xlabel('Cosine Similarity', color='#8b9ab0')\nax2.set_ylabel('Count', color='#8b9ab0')\nax2.tick_params(colors='#8b9ab0')\nfor s in ax2.spines.values(): s.set_color('#1a2030')\nax2.legend(facecolor='#1a2030', labelcolor='white')\n\nplt.suptitle('Jaguar Re-ID — Training & Inference Results',\n             color='#d4a853', fontsize=15, fontweight='bold')\nplt.tight_layout()\nplt.savefig(OUT_DIR/'results.png', dpi=150, bbox_inches='tight', facecolor='#0a0e17')\nplt.show()\nprint(\"Results visualization saved!\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"md_umap","cell_type":"markdown","source":"## 🗺️ Step 12 — UMAP Embedding Visualization","metadata":{}},{"id":"umap","cell_type":"code","source":"try:\n    import umap as umap_lib\n    m = model.module if N_GPU > 1 else model\n    m.eval()\n    all_e, all_l, all_n = [], [], []\n    max_s = 400 if not AMP else 800\n\n    with torch.no_grad():\n        for i in range(0, min(max_s, len(train_df)), 16):\n            batch = train_df.iloc[i:i+16]\n            imgs  = []\n            for _, row in batch.iterrows():\n                img = np.array(Image.open(TRAIN_DIR/row['filename']).convert('RGB'))\n                imgs.append(val_tfm(image=img)['image'])\n            imgs = torch.stack(imgs).to(DEVICE)\n            with autocast(enabled=AMP):\n                e = m(imgs).cpu().numpy()\n            all_e.append(e)\n            all_l.extend(batch['label'].tolist())\n            all_n.extend(batch['ground_truth'].tolist())\n\n    E_u = np.concatenate(all_e)\n    L_u = np.array(all_l)\n    print(f\"Computing UMAP on {len(E_u)} embeddings...\")\n\n    reducer = umap_lib.UMAP(n_components=2, random_state=42,\n                             n_neighbors=15, min_dist=0.1, metric='cosine')\n    E2d = reducer.fit_transform(E_u)\n\n    fig, ax = plt.subplots(figsize=(15, 11), facecolor='#0a0e17')\n    ax.set_facecolor('#0d1117')\n    cmap = plt.cm.get_cmap('tab20', NUM_CLS)\n\n    for uid in np.unique(L_u):\n        mask = L_u == uid\n        name = [n for n, l in zip(all_n, L_u) if l == uid][0]\n        ax.scatter(E2d[mask,0], E2d[mask,1], c=[cmap(uid/NUM_CLS)],\n                   s=30, alpha=0.85, edgecolors='none')\n        cx, cy = E2d[mask,0].mean(), E2d[mask,1].mean()\n        ax.annotate(name, (cx, cy), fontsize=6.5, color='white', alpha=0.95,\n                    ha='center', fontweight='bold',\n                    bbox=dict(boxstyle='round,pad=0.2', facecolor='#0a0e17',\n                              alpha=0.55, edgecolor='none'))\n\n    ax.set_title('UMAP — Jaguar Identity Embedding Space',\n                 color='#d4a853', fontsize=15, fontweight='bold', pad=15)\n    ax.set_xlabel('UMAP Dimension 1', color='#8b9ab0', fontsize=11)\n    ax.set_ylabel('UMAP Dimension 2', color='#8b9ab0', fontsize=11)\n    ax.tick_params(colors='#8b9ab0')\n    for s in ax.spines.values(): s.set_color('#1a2030')\n    plt.tight_layout()\n    plt.savefig(OUT_DIR/'umap.png', dpi=150, bbox_inches='tight', facecolor='#0a0e17')\n    plt.show()\n    print(\"UMAP saved!\")\nexcept Exception as e:\n    print(f\"UMAP skipped: {e}\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"md_heatmap","cell_type":"markdown","source":"## 🔥 Step 13 — Inter-Jaguar Similarity Heatmap","metadata":{}},{"id":"heatmap","cell_type":"code","source":"m = model.module if N_GPU > 1 else model\nm.eval()\ntop12 = train_df['ground_truth'].value_counts().head(12).index.tolist()\ndf12  = train_df[train_df['ground_truth'].isin(top12)]\n\nrep_embs, rep_names = {}, []\nwith torch.no_grad():\n    for name_j in top12:\n        rows = df12[df12['ground_truth'] == name_j].head(6)\n        imgs = []\n        for _, row in rows.iterrows():\n            img = np.array(Image.open(TRAIN_DIR/row['filename']).convert('RGB'))\n            imgs.append(val_tfm(image=img)['image'])\n        imgs = torch.stack(imgs).to(DEVICE)\n        with autocast(enabled=AMP):\n            emb = m(imgs).mean(dim=0)\n        rep_embs[name_j] = emb.cpu()\n        rep_names.append(name_j)\n\nemb_mat = torch.stack([rep_embs[n] for n in rep_names])\nsim_mat  = (emb_mat @ emb_mat.T).numpy()\n\nfig, ax = plt.subplots(figsize=(12, 10), facecolor='#0a0e17')\nsns.heatmap(sim_mat, annot=True, fmt='.2f', cmap='YlOrBr',\n            xticklabels=rep_names, yticklabels=rep_names,\n            ax=ax, linewidths=0.4, linecolor='#0a0e17',\n            annot_kws={'size': 8, 'color': 'black'})\nax.set_title('Inter-Individual Cosine Similarity Matrix\\n(Top 12 Jaguars — Average Embeddings)',\n             color='#d4a853', fontsize=13, fontweight='bold', pad=12)\nax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha='right', color='white', fontsize=9)\nax.set_yticklabels(ax.get_yticklabels(), rotation=0, color='white', fontsize=9)\nax.collections[0].colorbar.ax.tick_params(colors='white')\nax.collections[0].colorbar.ax.yaxis.label.set_color('white')\nplt.tight_layout()\nplt.savefig(OUT_DIR/'similarity_heatmap.png', dpi=150, bbox_inches='tight', facecolor='#0a0e17')\nplt.show()\nprint(\"Heatmap saved!\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"md_ap","cell_type":"markdown","source":"## 📐 Step 14 — Per-Identity AP & Embedding Separability","metadata":{}},{"id":"ap","cell_type":"code","source":"m = model.module if N_GPU > 1 else model\nm.eval()\nall_e2, all_l2 = [], []\n\nwith torch.no_grad():\n    for i in range(0, len(train_df), 32):\n        batch = train_df.iloc[i:i+32]\n        imgs  = []\n        for _, row in batch.iterrows():\n            img = np.array(Image.open(TRAIN_DIR/row['filename']).convert('RGB'))\n            imgs.append(val_tfm(image=img)['image'])\n        imgs = torch.stack(imgs).to(DEVICE)\n        with autocast(enabled=AMP):\n            e = m(imgs).cpu()\n        all_e2.append(e)\n        all_l2.extend(batch['label'].tolist())\n\nE_all  = torch.cat(all_e2).numpy()\nL_all  = np.array(all_l2)\nS_all  = E_all @ E_all.T\nlnames = {v: k for k, v in l2i.items()}\n\nap_scores = {}\nfor uid in range(NUM_CLS):\n    mask = L_all == uid\n    if mask.sum() < 2:\n        continue\n    aps = []\n    for qi in np.where(mask)[0]:\n        sims     = S_all[qi].copy()\n        sims[qi] = -999\n        is_match = (L_all == uid).astype(float)\n        is_match[qi] = 0\n        order    = np.argsort(-sims)\n        prec, nc = [], 0\n        for i, c in enumerate(is_match[order]):\n            if c:\n                nc += 1\n                prec.append(nc / (i+1))\n        aps.append(np.mean(prec) if prec else 0)\n    ap_scores[lnames[uid]] = np.mean(aps)\n\nsorted_ap = sorted(ap_scores.items(), key=lambda x: x[1], reverse=True)\nnames_ap  = [x[0] for x in sorted_ap]\nvals_ap   = [x[1] for x in sorted_ap]\nmean_ap   = np.mean(vals_ap)\n\nfig, axes = plt.subplots(1, 3, figsize=(22, 7), facecolor='#0a0e17')\n\nax = axes[0]\nax.set_facecolor('#0f1520')\nbar_colors = ['#d4a853' if v >= mean_ap else '#4ecdc4' for v in vals_ap]\nax.barh(names_ap[::-1], vals_ap[::-1], color=bar_colors[::-1],\n        edgecolor='#0a0e17', height=0.7)\nax.axvline(mean_ap, color='#ff6b35', linestyle='--', linewidth=2,\n           label=f'Mean AP = {mean_ap:.3f}')\nax.set_title('Per-Identity Average Precision', color='#d4a853', fontsize=12, fontweight='bold')\nax.set_xlabel('AP', color='#8b9ab0')\nax.tick_params(colors='#8b9ab0', labelsize=8)\nfor s in ax.spines.values(): s.set_color('#1a2030')\nax.legend(facecolor='#1a2030', labelcolor='white')\n\nintra_mask = np.array([[L_all[i]==L_all[j] and i!=j for j in range(len(L_all))] for i in range(len(L_all))])\ninter_mask = np.array([[L_all[i]!=L_all[j] for j in range(len(L_all))] for i in range(len(L_all))])\nintra_s    = S_all[intra_mask]\ninter_s    = S_all[inter_mask]\ninter_s    = np.random.choice(inter_s, min(len(intra_s)*3, len(inter_s)), replace=False)\n\nax2 = axes[1]\nax2.set_facecolor('#0f1520')\nax2.hist(intra_s, bins=50, alpha=0.75, color='#d4a853', density=True,\n         label=f'Same ID  mean={intra_s.mean():.3f}')\nax2.hist(inter_s, bins=50, alpha=0.75, color='#4ecdc4', density=True,\n         label=f'Diff ID  mean={inter_s.mean():.3f}')\ngap = intra_s.mean() - inter_s.mean()\nax2.set_title(f'Intra vs Inter-Class Similarity\\nSeparability gap = {gap:.3f}',\n              color='#d4a853', fontsize=12, fontweight='bold')\nax2.set_xlabel('Cosine Similarity', color='#8b9ab0')\nax2.set_ylabel('Density', color='#8b9ab0')\nax2.tick_params(colors='#8b9ab0')\nfor s in ax2.spines.values(): s.set_color('#1a2030')\nax2.legend(facecolor='#1a2030', labelcolor='white')\n\nax3 = axes[2]\nax3.set_facecolor('#0f1520')\nax3.axis('off')\nstats = [\n    ['Mean AP (proxy)', f'{mean_ap:.4f}'],\n    ['Best jaguar', f'{names_ap[0]}'],\n    ['Best AP', f'{vals_ap[0]:.4f}'],\n    ['Worst jaguar', f'{names_ap[-1]}'],\n    ['Worst AP', f'{vals_ap[-1]:.4f}'],\n    ['Intra-class mean', f'{intra_s.mean():.4f}'],\n    ['Inter-class mean', f'{inter_s.mean():.4f}'],\n    ['Separability gap', f'{gap:.4f}'],\n]\nt = ax3.table(cellText=stats, colLabels=['Metric', 'Value'],\n              cellLoc='center', loc='center',\n              cellColours=[['#1a2030','#1a2030']]*len(stats),\n              colColours=['#d4a853','#d4a853'])\nt.auto_set_font_size(False)\nt.set_fontsize(10)\nt.scale(1.3, 1.8)\nfor (r,c), cell in t.get_celld().items():\n    cell.set_text_props(color='white' if r>0 else '#0a0e17', fontweight='bold' if r==0 else 'normal')\n    cell.set_edgecolor('#1a2030')\nax3.set_title('Embedding Quality Metrics', color='#d4a853', fontsize=12, fontweight='bold', pad=10)\n\nplt.suptitle(f'Identity-Balanced mAP (proxy) = {mean_ap:.4f}',\n             color='#d4a853', fontsize=15, fontweight='bold')\nplt.tight_layout()\nplt.savefig(OUT_DIR/'per_class_ap.png', dpi=150, bbox_inches='tight', facecolor='#0a0e17')\nplt.show()\nprint(f\"Mean AP  : {mean_ap:.4f}\")\nprint(f\"Gap      : {gap:.4f}\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"md_preds","cell_type":"markdown","source":"## 🐆 Step 15 — Sample Predictions","metadata":{}},{"id":"predictions","cell_type":"code","source":"n_q     = min(6, len(all_files))\nq_files = all_files[:n_q]\nE_np    = test_embs.numpy()\n\nfig, axes = plt.subplots(n_q, 4, figsize=(16, n_q*2.8), facecolor='#0a0e17')\nif n_q == 1:\n    axes = axes.reshape(1, -1)\nfig.suptitle('Query Image  →  Top-3 Most Similar Gallery Images',\n             color='#d4a853', fontsize=14, fontweight='bold', y=1.01)\n\nfor row_i, qf in enumerate(q_files):\n    qi      = file2idx[qf]\n    gallery = [(f, float(np.dot(E_np[qi], E_np[file2idx[f]])))\n               for f in all_files if f != qf]\n    top3    = sorted(gallery, key=lambda x: x[1], reverse=True)[:3]\n\n    q_img = Image.open(TEST_DIR/qf).convert('RGB').resize((180, 180))\n    axes[row_i, 0].imshow(q_img)\n    axes[row_i, 0].set_title('QUERY', color='#4ecdc4', fontsize=9, fontweight='bold')\n    axes[row_i, 0].axis('off')\n\n    for col_j, (gf, sim) in enumerate(top3):\n        g_img = Image.open(TEST_DIR/gf).convert('RGB').resize((180, 180))\n        color = '#d4a853' if sim > 0.7 else '#ff6b35' if sim > 0.5 else '#8b9ab0'\n        axes[row_i, col_j+1].imshow(g_img)\n        axes[row_i, col_j+1].set_title(f'Rank {col_j+1} | {sim:.3f}',\n                                        color=color, fontsize=8, fontweight='bold')\n        axes[row_i, col_j+1].axis('off')\n\nfor ax in axes.flat:\n    ax.set_facecolor('#0a0e17')\n\nplt.tight_layout()\nplt.savefig(OUT_DIR/'sample_predictions.png', dpi=130, bbox_inches='tight', facecolor='#0a0e17')\nplt.show()\nprint(\"Sample predictions saved!\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"md_summary","cell_type":"markdown","source":"## 🏆 Step 16 — Final Summary","metadata":{}},{"id":"summary","cell_type":"code","source":"print(\"=\" * 58)\nprint(\"   JAGUAR RE-ID — PIPELINE COMPLETE\")\nprint(\"=\" * 58)\nprint(f\"  Device         : {DEVICE} ({N_GPU} GPU(s))\")\nprint(f\"  Backbone       : {BACKBONE}\")\nprint(f\"  Image size     : {IMG_SIZE}x{IMG_SIZE}\")\nprint(f\"  Epochs trained : {EPOCHS}\")\nprint(f\"  Best loss      : {best_loss:.4f}\")\nprint(f\"  Test pairs     : {len(sub_df):,}\")\nprint(f\"  Score range    : [{sub_df['score'].min():.4f}, {sub_df['score'].max():.4f}]\")\nprint(f\"  Score mean     : {sub_df['score'].mean():.4f}\")\nprint(\"-\" * 58)\nprint(\"  Output files:\")\nfor f in sorted(OUT_DIR.glob('*')):\n    sz = f.stat().st_size / 1024\n    icon = 'CSV' if f.suffix=='.csv' else 'PTH' if f.suffix=='.pth' else 'PNG'\n    print(f\"    [{icon}] {f.name:<35} {sz:>7.1f} KB\")\nprint(\"=\" * 58)\nprint(\"  Submission ready — submit submission.csv!\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}