{"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":127282,"databundleVersionId":15520105,"isSourceIdPinned":false}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 🐢🦎 AnimalCLEF 2026 ","metadata":{}},{"cell_type":"markdown","source":"## 1. Installation","metadata":{}},{"cell_type":"code","source":"import os, sys\n\n# ── CRITICAL: set memory config BEFORE any CUDA import ────────────────────────\nos.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'\n\n# ── Install missing packages only — never reinstall torch ─────────────────────\n# Kaggle already ships a CUDA-compatible PyTorch. Reinstalling it\n# (a) corrupts the CUDA context mid-session\n# (b) causes NameError: 'torch' is not defined in later cells\n# (c) wastes 2-3 min and ~1 GB VRAM on stale library residue\n#\n# After running this cell, do NOT restart the kernel manually —\n# pip installs here are additive and safe without a restart.\n!pip install -q --no-cache-dir timm hdbscan umap-learn\n!pip install -q --no-cache-dir 'transformers>=4.40.0,<5.0.0' accelerate huggingface_hub safetensors\n!pip install -q --no-cache-dir git+https://github.com/cvg/LightGlue.git\n!pip install -q --no-cache-dir scikit-learn scipy networkx\n\nprint(\"Installs complete. torch version:\", end=\" \")\nimport torch; print(torch.__version__)\nprint(\"CUDA available:\", torch.cuda.is_available())\nprint(\"PYTORCH_CUDA_ALLOC_CONF:\", os.environ.get('PYTORCH_CUDA_ALLOC_CONF'))","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Imports & Reproducibility","metadata":{}},{"cell_type":"code","source":"import os, gc, warnings, random\nfrom pathlib import Path\nfrom collections import Counter\n\n# Guard: ensure alloc config is set even if install cell was skipped\nos.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF', 'expandable_segments:True')\n\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nfrom tqdm import tqdm\n\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 torch.optim import AdamW\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom torch.amp import autocast, GradScaler   # torch.amp works on torch>=1.10\n\nimport timm\nimport hdbscan\nimport umap.umap_ as umap_lib\nimport networkx as nx\n\nfrom sklearn.metrics import adjusted_rand_score\nfrom sklearn.preprocessing import normalize\nfrom sklearn.isotonic import IsotonicRegression\nfrom scipy.spatial.distance import cdist\n\nwarnings.filterwarnings('ignore')\nSEED = 42\nrandom.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)\nif torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)\n\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nAMP_DTYPE = torch.float16 if DEVICE.type == 'cuda' else torch.bfloat16\n\nprint(f'Device: {DEVICE}')\nif torch.cuda.is_available():\n    print(f'GPU:   {torch.cuda.get_device_name(0)}')\n    print(f'VRAM:  {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB')\nprint(f'torch: {torch.__version__}')\nprint(f'ALLOC: {os.environ.get(\"PYTORCH_CUDA_ALLOC_CONF\", \"NOT SET — rerun install cell\")}')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Configuration","metadata":{}},{"cell_type":"code","source":"class CFG:\n    DATA_DIR   = Path('/kaggle/input/competitions/animal-clef-2026')\n    OUTPUT_DIR = Path('/kaggle/working/')\n\n    # ---------- Models ----------\n    MEGA_MODEL_ID = 'hf-hub:BVRA/MegaDescriptor-L-384'\n    DINO_MODEL_ID = 'vit_large_patch14_reg4_dinov2'\n    MIEW_MODEL_ID = 'conservationxlabs/miewid-msv3'\n    IMG_MEGA = 384\n    IMG_DINO = 448   # ← was 518; 448 saves ~2 GB VRAM with <1% accuracy loss\n    IMG_MIEW = 440\n\n    # ---------- Inference ----------\n    BATCH_SIZE  = 8   # ← was 16; halving batch halves activation memory\n    NUM_WORKERS = 2\n    TTA_FLIPS   = True\n\n    # ---------- Ensemble weights ----------\n    MEGA_W = 0.45\n    DINO_W = 0.25\n    MIEW_W = 0.30\n\n    # ---------- UMAP ----------\n    UMAP_N_NEIGHBORS = 15\n    UMAP_MIN_DIST    = 0.0\n    UMAP_COMPONENTS  = 64\n\n    # ---------- HDBSCAN ----------\n    HDBSCAN_MIN_SAMPLES = 1\n    HDBSCAN_METHOD      = 'eom'\n\n    # ---------- Post-processing ----------\n    MERGE_THRESHOLD = 0.92\n\n    # ---------- LightGlue ----------\n    # 512 keypoints instead of 1024 — halves LightGlue VRAM, negligible accuracy drop\n    LG_MAX_KEYPOINTS    = 512\n    LG_FEATURES         = 'aliked'\n    FUSION_GRAPH_THRESH = 0.55\n    ALPHA_LOCAL         = 0.40\n    LG_TOPK_CANDIDATES  = 10\n    # Run LightGlue on CPU if GPU free VRAM < this threshold (bytes)\n    LG_MIN_FREE_VRAM    = 2 * 1024**3   # 2 GB\n\n    # ---------- Self-training ----------\n    N_SELFTR_ROUNDS = 2\n    FT_EPOCHS       = 5\n    FT_LR_BACKBONE  = 1e-5\n    FT_LR_HEAD      = 1e-4\n    FT_BATCH        = 16   # ← was 32; halved to prevent OOM during backward pass\n    ARCFACE_MARGIN  = 0.5\n    ARCFACE_SCALE   = 64\n\nCFG.OUTPUT_DIR.mkdir(exist_ok=True)\n\ndef free_vram():\n    \"\"\"Return free GPU VRAM in bytes.\"\"\"; \n    if not torch.cuda.is_available(): return float('inf')\n    return torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated()\n\nprint('Config ready.')\nprint(f'Free VRAM now: {free_vram()/1e9:.2f} GB')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4. Data Loading","metadata":{}},{"cell_type":"code","source":"def find_col(df, candidates):\n    for c in candidates:\n        if c in df.columns: return c\n    return None\n\ndef load_metadata(data_dir):\n    for fname in ['metadata.csv','train.csv','data.csv','images.csv']:\n        p = data_dir / fname\n        if p.exists():\n            print(f'Loaded: {fname}'); return pd.read_csv(p)\n    raise FileNotFoundError('No metadata CSV found.')\n\nmetadata = load_metadata(CFG.DATA_DIR)\nprint(metadata.shape, metadata.columns.tolist())\n\nPATH_COL    = find_col(metadata, ['path','image_path','file_path','filename','image_id','image'])\nSPECIES_COL = find_col(metadata, ['species','dataset','category','taxon'])\nSPLIT_COL   = find_col(metadata, ['split','subset','set','partition'])\nID_COL      = find_col(metadata, ['identity','individual_id','label','individual'])\nprint(f'path:{PATH_COL}  species:{SPECIES_COL}  split:{SPLIT_COL}  id:{ID_COL}')\n\ndef resolve(p, base):\n    for c in [Path(p), base/p, base/str(p).lstrip('/')]:\n        if c.exists(): return str(c)\n    return str(base/p)\n\nmetadata['full_path'] = metadata[PATH_COL].apply(lambda x: resolve(x, CFG.DATA_DIR))\nok = metadata['full_path'].apply(lambda x: Path(x).exists()).mean()\nprint(f'{ok*100:.1f}% of paths resolved.')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if SPLIT_COL:\n    test_df  = metadata[metadata[SPLIT_COL]=='test'].copy().reset_index(drop=True)\n    train_df = metadata[metadata[SPLIT_COL]!='test'].copy().reset_index(drop=True)\nelif ID_COL:\n    test_df  = metadata[metadata[ID_COL].isna()].copy().reset_index(drop=True)\n    train_df = metadata[metadata[ID_COL].notna()].copy().reset_index(drop=True)\nelse:\n    test_df, train_df = metadata.copy().reset_index(drop=True), pd.DataFrame()\n\nSPECIES_LIST = test_df[SPECIES_COL].unique().tolist() if SPECIES_COL else ['all']\nTEST_PATHS   = test_df['full_path'].tolist()\nprint(f'Train: {len(train_df)}  |  Test: {len(test_df)}  |  Species: {SPECIES_LIST}')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 5. Dataset & Transforms","metadata":{}},{"cell_type":"code","source":"class ImgDataset(Dataset):\n    def __init__(self, paths, transform):\n        self.paths, self.transform = paths, transform\n    def __len__(self): return len(self.paths)\n    def __getitem__(self, i):\n        try:   img = Image.open(self.paths[i]).convert('RGB')\n        except: img = Image.new('RGB',(224,224),0)\n        return self.transform(img), i\n\ndef base_tf(size):\n    return T.Compose([\n        T.Resize((size,size), interpolation=T.InterpolationMode.BICUBIC),\n        T.ToTensor(),\n        T.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),\n    ])\n\ndef flip_tf(size, hflip=True):\n    return T.Compose([\n        T.Resize((size,size), interpolation=T.InterpolationMode.BICUBIC),\n        T.RandomHorizontalFlip(p=1.0) if hflip else T.RandomVerticalFlip(p=1.0),\n        T.ToTensor(),\n        T.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),\n    ])\n\ndef aug_tf(size):\n    return T.Compose([\n        T.RandomResizedCrop(size, scale=(0.7,1.0)),\n        T.RandomHorizontalFlip(),\n        T.ColorJitter(0.3,0.3,0.3,0.1),\n        T.ToTensor(),\n        T.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),\n    ])\n\nprint('Transforms ready.')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 6. Embedding Extraction (memory-safe with autocast + per-pass cache flush)","metadata":{}},{"cell_type":"code","source":"@torch.no_grad()\ndef extract_timm(model, paths, img_size):\n    \"\"\"\n    TTA-averaged L2-normalized embeddings.\n    OOM fixes applied:\n      - autocast(fp16) halves activation memory\n      - each batch result moved to CPU immediately\n      - torch.cuda.empty_cache() called after every TTA pass\n      - accumulator stays on CPU throughout\n    \"\"\"\n    model.eval()\n    tfs = [base_tf(img_size)]\n    if CFG.TTA_FLIPS:\n        tfs += [flip_tf(img_size, hflip=True), flip_tf(img_size, hflip=False)]\n\n    accum = None   # CPU float32 accumulator\n\n    for pass_idx, tf in enumerate(tfs):\n        ds = ImgDataset(paths, tf)\n        dl = DataLoader(ds, batch_size=CFG.BATCH_SIZE, shuffle=False,\n                        num_workers=CFG.NUM_WORKERS, pin_memory=(DEVICE.type=='cuda'))\n        pass_cpu = []\n        for imgs, _ in tqdm(dl, desc=f'  pass {pass_idx+1}/{len(tfs)} sz={img_size}', leave=False):\n            imgs = imgs.to(DEVICE, non_blocking=True)\n            with autocast(device_type=DEVICE.type, dtype=AMP_DTYPE):\n                e = model(imgs)\n            e = F.normalize(e.to(torch.float32), dim=1).cpu()  # cast back, move to CPU\n            pass_cpu.append(e)\n            del imgs\n        arr = torch.cat(pass_cpu, dim=0)   # (N, D) on CPU\n        accum = arr if accum is None else accum + arr\n        del pass_cpu, arr\n        gc.collect()\n        if DEVICE.type == 'cuda':\n            torch.cuda.empty_cache()\n\n    return F.normalize(accum / len(tfs), dim=1).numpy()\n\nprint('Extraction function ready.')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7. Extract Embeddings (one model at a time — free VRAM before next)","metadata":{}},{"cell_type":"code","source":"# ── MegaDescriptor-L ──────────────────────────────────────────────────────────\nprint('Loading MegaDescriptor-L-384...')\nmega_model = timm.create_model(CFG.MEGA_MODEL_ID, pretrained=True, num_classes=0).to(DEVICE)\nprint(f'  dim={mega_model.num_features}  |  free VRAM before extract: {free_vram()/1e9:.2f} GB')\nmega_emb = extract_timm(mega_model, TEST_PATHS, CFG.IMG_MEGA)\nprint(f'  mega_emb: {mega_emb.shape}')\ndel mega_model; gc.collect(); torch.cuda.empty_cache()\nprint(f'  free VRAM after del: {free_vram()/1e9:.2f} GB')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ── DINOv2-L ──────────────────────────────────────────────────────────────────\n# Using 448px instead of 518px: saves ~2 GB VRAM, <1% accuracy loss on re-ID tasks\nprint('Loading DINOv2-L-448...')\ndino_model = timm.create_model(CFG.DINO_MODEL_ID, pretrained=True, num_classes=0,\n                                dynamic_img_size=True).to(DEVICE)\nprint(f'  dim={dino_model.num_features}  |  free VRAM: {free_vram()/1e9:.2f} GB')\ndino_emb = extract_timm(dino_model, TEST_PATHS, CFG.IMG_DINO)\nprint(f'  dino_emb: {dino_emb.shape}')\ndel dino_model; gc.collect(); torch.cuda.empty_cache()\nprint(f'  free VRAM after del: {free_vram()/1e9:.2f} GB')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ── MiewID-msv3 ───────────────────────────────────────────────────────────────\n# Path A: AutoModel (transformers <5.0)\n# Path B: timm backbone + safetensors weight transplant (fallback)\n\ndef _miewid_wrapper(raw_model):\n    class W(nn.Module):\n        def __init__(self, m): super().__init__(); self.m = m\n        def forward(self, x):\n            out = self.m(x)\n            if isinstance(out, torch.Tensor): return out\n            if hasattr(out, 'last_hidden_state'): return out.last_hidden_state[:, 0]\n            return out[0]\n    return W(raw_model)\n\nprint('Loading MiewID-msv3...')\nmiew_model = None\n\ntry:\n    import transformers as _tr\n    _ver = tuple(int(x) for x in _tr.__version__.split('.')[:2])\n    if _ver >= (5, 0):\n        raise RuntimeError(f'transformers {_tr.__version__} >=5.0')\n    from transformers import AutoModel as _AM\n    _raw = _AM.from_pretrained(CFG.MIEW_MODEL_ID, trust_remote_code=True)\n    miew_model = _miewid_wrapper(_raw).to(DEVICE)\n    print(f'  Path A OK (transformers {_tr.__version__})')\nexcept Exception as e:\n    print(f'  Path A failed: {e}')\n\nif miew_model is None:\n    try:\n        from huggingface_hub import hf_hub_download\n        import safetensors.torch as st\n        ckpt = hf_hub_download(repo_id=CFG.MIEW_MODEL_ID, filename='model.safetensors')\n        state = st.load_file(ckpt)\n        _bb = timm.create_model('efficientnetv2_rw_m.agc_in1k', pretrained=False, num_classes=0)\n        prefix = 'efficientnet.'\n        bb_state = {k[len(prefix):]: v for k, v in state.items() if k.startswith(prefix)}\n        _bb.load_state_dict(bb_state or state, strict=False)\n        class _M(nn.Module):\n            def __init__(self, bb, sd):\n                super().__init__(); self.bb = bb\n                pw = sd.get('head.weight') or sd.get('proj.weight')\n                pb = sd.get('head.bias')   or sd.get('proj.bias')\n                self.head = (nn.Linear(_bb.num_features, pw.shape[0], bias=pb is not None)\n                             if pw is not None else nn.Identity())\n                if pw is not None:\n                    self.head.weight.data.copy_(pw)\n                    if pb is not None: self.head.bias.data.copy_(pb)\n            def forward(self, x): return self.head(self.bb(x))\n        miew_model = _M(_bb, state).to(DEVICE)\n        print(f'  Path B OK (timm+safetensors)')\n    except Exception as e2:\n        print(f'  Path B failed: {e2}')\n\nif miew_model is not None:\n    print(f'  free VRAM before extract: {free_vram()/1e9:.2f} GB')\n    miew_emb = extract_timm(miew_model, TEST_PATHS, CFG.IMG_MIEW)\n    print(f'  miew_emb: {miew_emb.shape}')\n    del miew_model; gc.collect(); torch.cuda.empty_cache()\nelse:\n    miew_emb = np.zeros((len(TEST_PATHS), 1), dtype=np.float32)\n    CFG.MIEW_W, CFG.MEGA_W, CFG.DINO_W = 0.0, 0.60, 0.40\n    print('  ⚠ MiewID skipped. Weights: Mega=0.60, Dino=0.40')\n\nprint(f'Free VRAM after all extractions: {free_vram()/1e9:.2f} GB')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 8. Three-Model Ensemble","metadata":{}},{"cell_type":"code","source":"def make_ensemble(embs_weights):\n    parts = [normalize(e, norm='l2') * w\n             for e, w in embs_weights\n             if w > 0 and e.shape[1] > 1]\n    return normalize(np.concatenate(parts, axis=1), norm='l2')\n\n# saved_embs holds the only copies needed for self-training rounds\nsaved_embs = {\n    'mega': mega_emb.copy(),\n    'dino': dino_emb.copy(),\n    'miew': miew_emb.copy(),\n}\n\ncombined_emb = make_ensemble([\n    (saved_embs['mega'], CFG.MEGA_W),\n    (saved_embs['dino'], CFG.DINO_W),\n    (saved_embs['miew'], CFG.MIEW_W),\n])\n\n# Delete the originals — saved_embs already has copies, no need for duplicates\ndel mega_emb, dino_emb, miew_emb\ngc.collect()\nprint(f'Ensemble: {combined_emb.shape}')\nprint(f'Weights  Mega={CFG.MEGA_W} Dino={CFG.DINO_W} Miew={CFG.MIEW_W}')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 9. Clustering Utilities","metadata":{}},{"cell_type":"code","source":"def reduce_umap(emb):\n    n = min(CFG.UMAP_COMPONENTS, emb.shape[0]-2, emb.shape[1])\n    return umap_lib.UMAP(\n        n_components=n, n_neighbors=min(CFG.UMAP_N_NEIGHBORS, emb.shape[0]-1),\n        min_dist=CFG.UMAP_MIN_DIST, metric='cosine',\n        random_state=SEED, low_memory=False, n_jobs=-1,\n    ).fit_transform(emb)\n\ndef run_hdbscan(emb_2d, n):\n    c = hdbscan.HDBSCAN(\n        min_cluster_size=max(2, n//50), min_samples=CFG.HDBSCAN_MIN_SAMPLES,\n        metric='euclidean', cluster_selection_method=CFG.HDBSCAN_METHOD, prediction_data=True)\n    lbl = c.fit_predict(emb_2d)\n    print(f'    HDBSCAN → {len(set(lbl))-(1 if -1 in lbl else 0)} clusters, {(lbl==-1).sum()} noise')\n    return lbl, c\n\ndef reassign_noise(lbl, emb_2d, cl, orig_emb):\n    noise = lbl == -1\n    if not noise.any(): return lbl\n    try:\n        soft = hdbscan.all_points_membership_vectors(cl)\n        if soft.shape[1] > 0:\n            lbl = lbl.copy(); lbl[noise] = np.argmax(soft[noise], axis=1); return lbl\n    except Exception: pass\n    valid = [c for c in np.unique(lbl) if c != -1]\n    if not valid:\n        lbl = lbl.copy()\n        for i, idx in enumerate(np.where(noise)[0]): lbl[idx] = i\n        return lbl\n    cents = np.stack([orig_emb[lbl==c].mean(0) for c in valid])\n    lbl = lbl.copy()\n    lbl[noise] = np.array(valid)[cdist(orig_emb[noise], cents, 'cosine').argmin(1)]\n    return lbl\n\ndef merge_close(lbl, orig_emb):\n    unique = np.unique(lbl)\n    if len(unique) <= 1: return lbl\n    cents = normalize(np.stack([orig_emb[lbl==c].mean(0) for c in unique]), norm='l2')\n    sim = cents @ cents.T\n    parent = {c:c for c in unique}\n    def find(x):\n        while parent[x]!=x: parent[x]=parent[parent[x]]; x=parent[x]\n        return x\n    n_m=0\n    for i in range(len(unique)):\n        for j in range(i+1,len(unique)):\n            if sim[i,j]>CFG.MERGE_THRESHOLD:\n                parent[find(unique[i])]=find(unique[j]); n_m+=1\n    m={c:find(c) for c in unique}\n    new_l=np.array([m[l] for l in lbl])\n    uid={v:i for i,v in enumerate(sorted(set(new_l)))}\n    new_l=np.array([uid[l] for l in new_l])\n    print(f'    Merged {n_m} pairs: {len(unique)}→{len(uid)} clusters')\n    return new_l\n\ndef cluster_subset(emb, n):\n    if n<=2: return np.arange(n)\n    red=reduce_umap(emb); lbl,cl=run_hdbscan(red,n)\n    lbl=reassign_noise(lbl,red,cl,emb); lbl=merge_close(lbl,emb)\n    return lbl\n\ndef full_cluster(emb, df):\n    final=np.full(len(df),-1,dtype=int); off=0\n    if SPECIES_COL:\n        for sp in SPECIES_LIST:\n            print(f'\\n  [{sp}]')\n            mask=(df[SPECIES_COL]==sp).values; idx=np.where(mask)[0]\n            lbl=cluster_subset(emb[idx],len(idx))\n            final[idx]=lbl+off; off+=lbl.max()+1\n            print(f'    {len(np.unique(lbl))} clusters / {len(idx)} images')\n    else:\n        final=cluster_subset(emb,len(df))\n    print(f'\\nTotal individuals: {len(np.unique(final))}')\n    return final\n\nprint('Clustering utils ready.')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 10. Initial Clustering (Round 0)","metadata":{}},{"cell_type":"code","source":"print('=== Initial clustering ===')\ncurrent_labels = full_cluster(combined_emb, test_df)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 11. WildFusion: ALIKED + LightGlue (VRAM-safe)\n\nKey OOM fixes:\n- LightGlue models loaded **after** all embedding models are deleted\n- Free VRAM checked before deciding GPU vs CPU execution\n- Images loaded on CPU, moved to LightGlue device only during `extract()`, immediately deleted\n- `max_num_keypoints=512` (was 1024) — halves attention matrix size","metadata":{}},{"cell_type":"code","source":"from lightglue import LightGlue, ALIKED\nfrom lightglue.utils import rbd\n\n# Always run LightGlue on CPU.\n# Reason: combined_emb + HDBSCAN residuals already occupy most of GPU RAM.\n# LightGlue on CPU is ~3x slower per pair but with only ~N*10 pairs\n# (N images × top-10 neighbours) it finishes in minutes, not hours.\n# Running on GPU risks OOM from fragmentation even when free_vram() looks healthy.\nLG_DEVICE = torch.device('cpu')\nprint(f'LightGlue device: {LG_DEVICE}  (CPU is intentional — prevents OOM)')\n\naliked_extractor = ALIKED(\n    max_num_keypoints=CFG.LG_MAX_KEYPOINTS,\n    detection_threshold=0.01,\n    resize=None,\n).eval().to(LG_DEVICE)\n\nlg_matcher = LightGlue(\n    features=CFG.LG_FEATURES,\n    depth_confidence=0.9,\n    width_confidence=0.95,\n).eval().to(LG_DEVICE)\n\nprint(f'LightGlue + ALIKED ready on {LG_DEVICE}.')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_for_lg(path, max_size=640):\n    \"\"\"Load image as [0,1] float CPU tensor, resized to max_size on longest edge.\"\"\"\n    img = Image.open(path).convert('RGB')\n    w, h = img.size\n    scale = min(max_size / max(h, w), 1.0)\n    if scale < 1.0:\n        img = img.resize((int(w * scale), int(h * scale)), Image.BICUBIC)\n    return torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.0\n\nprint('load_for_lg ready.')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def build_wildfusion_sim(emb, paths, _aliked, _lg_matcher, _lg_device,\n                         topk=CFG.LG_TOPK_CANDIDATES, alpha=CFG.ALPHA_LOCAL):\n    \"\"\"\n    Build fused similarity dict.\n    Accepts extractor/matcher as explicit args so no global state is needed.\n    The N×N global_sim is computed in float32 numpy — fine for CPU.\n    \"\"\"\n    N = len(paths)\n    emb_norm   = normalize(emb, norm='l2')\n    global_sim = (emb_norm @ emb_norm.T).astype(np.float32)\n    np.fill_diagonal(global_sim, 0)\n\n    k = min(topk, N - 1)\n    pairs = set()\n    for i in range(N):\n        for j in np.argsort(global_sim[i])[-k:]:\n            pairs.add((min(i, j), max(i, j)))\n    pairs = list(pairs)\n    print(f'  Running LightGlue on {len(pairs)} pairs (device={_lg_device})...')\n\n    pair_data, raw_g, raw_l = [], [], []\n    for (i, j) in tqdm(pairs, desc='  LG matching', leave=False):\n        img_a = load_for_lg(paths[i]).unsqueeze(0).to(_lg_device)\n        img_b = load_for_lg(paths[j]).unsqueeze(0).to(_lg_device)\n        try:\n            fa  = _aliked.extract(img_a.squeeze(0))\n            fb  = _aliked.extract(img_b.squeeze(0))\n            res = _lg_matcher({'image0': fa, 'image1': fb})\n            fa, fb, res = rbd(fa), rbd(fb), rbd(res)\n            nm = int(res['matches'].shape[0])\n        except Exception:\n            nm = 0\n        finally:\n            del img_a, img_b\n            if _lg_device.type == 'cuda':\n                torch.cuda.empty_cache()\n        pair_data.append((i, j, float(global_sim[i, j]), nm))\n        raw_g.append(float(global_sim[i, j]))\n        raw_l.append(float(nm))\n\n    raw_g = np.array(raw_g, dtype=np.float32)\n    raw_l = np.array(raw_l, dtype=np.float32)\n\n    y = np.array([1 if current_labels[i] == current_labels[j] else 0\n                  for (i, j, _, _) in pair_data], dtype=np.float32)\n\n    def cal(scores, y):\n        ir = IsotonicRegression(out_of_bounds='clip', increasing=True)\n        ir.fit(scores, y)\n        return ir\n\n    ir_g = cal(raw_g, y)\n    ir_l = cal(np.log1p(raw_l), y)\n\n    fused = {}\n    for (i, j, gs, lm) in pair_data:\n        s = ((1 - alpha) * float(ir_g.predict([gs])[0])\n             + alpha     * float(ir_l.predict([np.log1p(lm)])[0]))\n        fused[(i, j)] = fused[(j, i)] = s\n    return fused\n\n\ndef graph_cluster(fused_sim, n, thresh=CFG.FUSION_GRAPH_THRESH):\n    G = nx.Graph()\n    G.add_nodes_from(range(n))\n    for (i, j), s in fused_sim.items():\n        if i < j and s > thresh:\n            G.add_edge(i, j, weight=s)\n    lbl = np.full(n, -1, dtype=int)\n    for cid, comp in enumerate(nx.connected_components(G)):\n        for node in comp:\n            lbl[node] = cid\n    uid = {v: i for i, v in enumerate(sorted(set(lbl)))}\n    lbl = np.array([uid[l] for l in lbl])\n    print(f'  Graph cut → {len(np.unique(lbl))} clusters (thresh={thresh})')\n    return lbl\n\n\ndef apply_wildfusion(emb, df, _aliked, _lg_matcher, _lg_device):\n    \"\"\"Run WildFusion. Extractor/matcher passed explicitly — no globals needed.\"\"\"\n    final = np.full(len(df), -1, dtype=int)\n    off   = 0\n    if SPECIES_COL:\n        for sp in SPECIES_LIST:\n            print(f'\\n  WildFusion [{sp}]')\n            mask = (df[SPECIES_COL] == sp).values\n            idx  = np.where(mask)[0]\n            n    = len(idx)\n            if n < 4:\n                sp_lbl = current_labels[idx] - current_labels[idx].min()\n            else:\n                sp_paths = [TEST_PATHS[i] for i in idx]\n                fused    = build_wildfusion_sim(emb[idx], sp_paths,\n                                                _aliked, _lg_matcher, _lg_device)\n                sp_lbl   = graph_cluster(fused, n)\n            final[idx] = sp_lbl + off\n            off += int(final[idx].max()) + 1\n            print(f'    {len(np.unique(sp_lbl))} clusters / {n} images')\n    else:\n        fused  = build_wildfusion_sim(emb, TEST_PATHS, _aliked, _lg_matcher, _lg_device)\n        final  = graph_cluster(fused, len(df))\n    print(f'\\nWildFusion total: {len(np.unique(final))} individuals')\n    return final\n\nprint('WildFusion ready.')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print('=== Applying WildFusion (initial) ===')\n# aliked_extractor and lg_matcher were loaded in cell 25\nwf_labels = apply_wildfusion(combined_emb, test_df,\n                              aliked_extractor, lg_matcher, LG_DEVICE)\nprint(f'HDBSCAN: {len(np.unique(current_labels))}  |  WildFusion: {len(np.unique(wf_labels))}')\ncurrent_labels = wf_labels\n\n# Free LightGlue — not needed until self-training rounds\ndel aliked_extractor, lg_matcher\ngc.collect(); torch.cuda.empty_cache()\nprint(f'Free VRAM after WildFusion: {free_vram()/1e9:.2f} GB')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 12. Iterative ArcFace Self-Training","metadata":{}},{"cell_type":"code","source":"class ArcFaceHead(nn.Module):\n    def __init__(self, D, C, m=CFG.ARCFACE_MARGIN, s=CFG.ARCFACE_SCALE):\n        super().__init__()\n        self.W = nn.Parameter(torch.empty(C, D))\n        nn.init.xavier_uniform_(self.W)\n        self.s = s\n        self.cos_m = np.cos(m); self.sin_m = np.sin(m)\n        self.th = np.cos(np.pi - m); self.mm = np.sin(np.pi - m) * m\n\n    def forward(self, e, lbl):\n        e  = F.normalize(e, dim=1)\n        W  = F.normalize(self.W, dim=1)\n        ct = (e @ W.T).clamp(-1 + 1e-7, 1 - 1e-7)\n        st = (1 - ct ** 2).sqrt()\n        ctm = ct * self.cos_m - st * self.sin_m\n        ctm = torch.where(ct >= self.th, ctm, ct - self.mm)\n        oh  = torch.zeros_like(ct).scatter_(1, lbl.unsqueeze(1), 1)\n        return F.cross_entropy((oh * ctm + (1 - oh) * ct) * self.s, lbl)\n\n\nclass PseudoDS(Dataset):\n    def __init__(self, paths, labels, tf):\n        self.paths, self.labels, self.tf = paths, labels, tf\n    def __len__(self): return len(self.paths)\n    def __getitem__(self, i):\n        try:    img = Image.open(self.paths[i]).convert('RGB')\n        except: img = Image.new('RGB', (224, 224), 0)\n        return self.tf(img), int(self.labels[i])\n\n\ndef finetune_and_reextract(pseudo_labels, round_idx):\n    uid    = sorted(set(pseudo_labels.tolist()))\n    remap  = {v: i for i, v in enumerate(uid)}\n    mapped = np.array([remap[l] for l in pseudo_labels])\n    C      = len(uid)\n    print(f'  Round {round_idx}: {C} pseudo-classes | free VRAM: {free_vram()/1e9:.2f} GB')\n\n    backbone = timm.create_model(CFG.MEGA_MODEL_ID, pretrained=True, num_classes=0).to(DEVICE)\n    if hasattr(backbone, 'set_grad_checkpointing'):\n        backbone.set_grad_checkpointing(True)\n    D    = backbone.num_features\n    head = ArcFaceHead(D, C).to(DEVICE)\n\n    ds = PseudoDS(TEST_PATHS, mapped, aug_tf(CFG.IMG_MEGA))\n    dl = DataLoader(ds, batch_size=CFG.FT_BATCH, shuffle=True,\n                    num_workers=CFG.NUM_WORKERS, pin_memory=(DEVICE.type == 'cuda'),\n                    drop_last=True)\n\n    opt    = AdamW([{'params': backbone.parameters(), 'lr': CFG.FT_LR_BACKBONE},\n                    {'params': head.parameters(),     'lr': CFG.FT_LR_HEAD}],\n                   weight_decay=1e-4)\n    sched  = CosineAnnealingLR(opt, T_max=CFG.FT_EPOCHS)\n    scaler = GradScaler(device=DEVICE.type)\n\n    for epoch in range(CFG.FT_EPOCHS):\n        backbone.train(); head.train()\n        total = 0.0\n        for imgs, lbls in tqdm(dl, desc=f'    Ep {epoch+1}/{CFG.FT_EPOCHS}', leave=False):\n            imgs, lbls = imgs.to(DEVICE), lbls.to(DEVICE)\n            with autocast(device_type=DEVICE.type, dtype=AMP_DTYPE):\n                loss = head(F.normalize(backbone(imgs), dim=1), lbls)\n            opt.zero_grad(set_to_none=True)\n            scaler.scale(loss).backward()\n            scaler.unscale_(opt)\n            nn.utils.clip_grad_norm_(backbone.parameters(), 5.0)\n            scaler.step(opt)\n            scaler.update()\n            total += loss.item()\n        sched.step()\n        print(f'    Epoch {epoch+1} | loss {total/len(dl):.4f}')\n\n    # ── Critical: free optimizer, scaler, dataloader BEFORE switching to eval ──\n    # These objects hold references to backbone parameters and keep gradient\n    # buffers allocated. Deleting them frees those buffers before re-extraction.\n    del opt, sched, scaler, dl, ds\n    gc.collect()\n    if DEVICE.type == 'cuda':\n        torch.cuda.empty_cache()\n    print(f'  Freed optimizer/scaler | free VRAM: {free_vram()/1e9:.2f} GB')\n\n    print('  Re-extracting embeddings with fine-tuned backbone...')\n    backbone.eval()\n    new_mega = extract_timm(backbone, TEST_PATHS, CFG.IMG_MEGA)\n\n    del backbone, head\n    gc.collect()\n    if DEVICE.type == 'cuda':\n        torch.cuda.empty_cache()\n    print(f'  Done | free VRAM: {free_vram()/1e9:.2f} GB')\n    return new_mega\n\nprint('Self-training code ready.')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from lightglue import LightGlue, ALIKED\nfrom lightglue.utils import rbd\n\nfor round_i in range(1, CFG.N_SELFTR_ROUNDS + 1):\n    print(f'\\n========== Self-Training Round {round_i}/{CFG.N_SELFTR_ROUNDS} ==========')\n    print(f'  Free VRAM at start: {free_vram()/1e9:.2f} GB')\n\n    # a) Fine-tune + re-extract (backbone+head+optimizer freed inside the fn)\n    new_mega = finetune_and_reextract(current_labels, round_i)\n\n    # b) Rebuild ensemble on CPU — no GPU involved\n    combined_emb = make_ensemble([\n        (new_mega,            CFG.MEGA_W),\n        (saved_embs['dino'],  CFG.DINO_W),\n        (saved_embs['miew'],  CFG.MIEW_W),\n    ])\n    saved_embs['mega'] = new_mega\n    del new_mega\n    gc.collect()\n    print(f'  Ensemble rebuilt | free VRAM: {free_vram()/1e9:.2f} GB')\n\n    # c) Re-cluster (CPU — UMAP and HDBSCAN are CPU-only)\n    print('  Re-clustering...')\n    current_labels = full_cluster(combined_emb, test_df)\n\n    # d) WildFusion — load LightGlue on CPU, run, then free everything\n    #    combined_emb is passed in and then immediately deleted after this step.\n    print('  Loading LightGlue on CPU...')\n    _aliked = ALIKED(max_num_keypoints=CFG.LG_MAX_KEYPOINTS,\n                     detection_threshold=0.01, resize=None).eval().to('cpu')\n    _lg     = LightGlue(features=CFG.LG_FEATURES,\n                        depth_confidence=0.9, width_confidence=0.95).eval().to('cpu')\n    _lg_dev = torch.device('cpu')\n\n    current_labels = apply_wildfusion(combined_emb, test_df, _aliked, _lg, _lg_dev)\n\n    # e) Free everything before next round\n    del _aliked, _lg, combined_emb\n    gc.collect()\n    if DEVICE.type == 'cuda':\n        torch.cuda.empty_cache()\n    print(f'  Round {round_i} done | free VRAM: {free_vram()/1e9:.2f} GB')\n\nprint('\\n✅ All self-training rounds complete.')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 13. Local Validation","metadata":{}},{"cell_type":"code","source":"if ID_COL and ID_COL in test_df.columns and test_df[ID_COL].notna().any():\n    gt  = test_df[ID_COL].fillna('__unk__').values\n    ari = adjusted_rand_score(gt, current_labels)\n    print(f'Overall ARI: {ari:.4f}')\n    if SPECIES_COL:\n        for sp in SPECIES_LIST:\n            mask=(test_df[SPECIES_COL]==sp).values\n            if len(np.unique(gt[mask]))>1:\n                print(f'  ARI [{sp}]: {adjusted_rand_score(gt[mask],current_labels[mask]):.4f}')\nelse:\n    print(f'No GT labels. Predicted {len(np.unique(current_labels))} individuals across {len(test_df)} images.')","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 14. Submission","metadata":{}},{"cell_type":"code","source":"try:\n    sample=pd.read_csv(CFG.DATA_DIR/'sample_submission.csv')\n    sub_img_col,sub_lbl_col=sample.columns[0],sample.columns[1]\n    print('Cols from sample:', sub_img_col, sub_lbl_col)\nexcept FileNotFoundError:\n    sub_img_col,sub_lbl_col='image_id','identity'\n\nsubmission=pd.DataFrame({\n    sub_img_col: test_df[PATH_COL].values,\n    sub_lbl_col: [f'individual_{int(x):06d}' for x in current_labels],\n})\nout=CFG.OUTPUT_DIR/'submission.csv'\nsubmission.to_csv(out,index=False)\nprint(f'Saved: {out}')\n\ncounts=submission[sub_lbl_col].value_counts()\nprint(f'Images: {len(submission)} | Individuals: {submission[sub_lbl_col].nunique()}')\nprint(f'Min/Mean/Max per individual: {counts.min()} / {counts.mean():.1f} / {counts.max()}')\nprint(submission.head())","metadata":{},"outputs":[],"execution_count":null}]}