{"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":"gpu","dataSources":[{"sourceType":"competition","sourceId":126777,"databundleVersionId":15314950,"isSourceIdPinned":false},{"sourceType":"kernelVersion","sourceId":302156577,"isSourceIdPinned":false},{"sourceType":"kernelVersion","sourceId":302177944,"isSourceIdPinned":false},{"sourceType":"kernelVersion","sourceId":302226826,"isSourceIdPinned":false}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nv07d: Pseudo-Labeling — Ensemble-based self-training for all 3 models\n=====================================================================\nPhase 2 of the v07 ensemble pipeline.\n\nStrategy:\n  1. Load pre-trained weights + embeddings from v07a, v07b, v07f\n  2. Average all 3 embeddings → compute cosine sim to class centroids\n     (all models share 1024-dim space → clean average, no concat noise)\n  3. Assign pseudo-labels where all 3 models agree AND confidence > 0.90\n  4. Fine-tune each model for 5 epochs at lower LR on train + pseudo data\n  5. Save updated weights + embeddings for the ensemble notebook (v07e)\n\nModels:\n  - v07a: EVA-02 Large (seed=42, 20 epochs)\n  - v07b: EVA-02 Large (seed=123, 20 epochs)\n  - v07f: DINOv2-Large-reg4 (self-supervised, 15 epochs)\n  All produce 1024-dim embeddings → averaged for centroid matching.\n\nInputs (uploaded as Kaggle datasets from v07a/b/f outputs):\n  - model_v07a.pth + embeddings + filenames\n  - model_v07b.pth + embeddings + filenames\n  - model_v07f.pth + embeddings + filenames\n\nOutputs to /kaggle/working/output/:\n  - model_v07a_pl.pth, model_v07b_pl.pth, model_v07f_pl.pth\n  - embeddings_v07{a,b,f}_pl_{train,test}.npy\n  - filenames_v07{a,b,f}_pl_{train,test}.json\n  - pseudo_labels.csv (for analysis)\n  - submission_v07d.csv (standalone from best single model after PL)\n\"\"\"\n\n# ── Install / imports ────────────────────────────────────────────────────\nimport subprocess, sys\nsubprocess.check_call([sys.executable, '-m', 'pip', 'install', '-qU', 'timm'])\n\nimport os, math, random, json, time\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom copy import deepcopy\nfrom tqdm.auto import tqdm\nfrom PIL import Image\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 transforms\nimport timm\n\n# ── Config ───────────────────────────────────────────────────────────────\n\nVERSION = 'v07d'\nSEED = 42\n\nNUM_CLASSES = 31\nPL_THRESHOLD = 0.90     # confidence threshold for pseudo-labeling\nPL_MAX_ADD = 500         # safety cap\n\n# Fine-tuning config (same for all 3 models)\nPL_EPOCHS = 5\nPL_LR = 1e-5            # lower LR for fine-tuning on pseudo data\nWEIGHT_DECAY = 1e-3\nEMA_DECAY = 0.999\n\nARCFACE_S = 30.0\nARCFACE_M = 0.5\n\nUSE_TTA = True\n\nNORM_MEAN = [0.481, 0.457, 0.408]\nNORM_STD = [0.268, 0.261, 0.275]\n\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nDEVICE_TYPE = 'cuda' if torch.cuda.is_available() else 'cpu'\n\n# ── Paths ────────────────────────────────────────────────────────────────\n\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(parents=True, exist_ok=True)\n\n# ── Auto-detect base model output directories ────────────────────────────\n# When you add a notebook's output as input on Kaggle, its files appear at:\n#   /kaggle/input/<notebook-url-slug>/output/\n# The slug is shown in the notebook URL, e.g. \"v07a-eva02-seed42\" → slug.\n# We search /kaggle/input/ for the marker files so you don't need to hardcode.\n\ndef find_model_dir(version: str) -> Path:\n    \"\"\"Search /kaggle/input/ recursively for the directory containing model_{version}.pth.\"\"\"\n    search_root = Path('/kaggle/input')\n    marker = f'model_{version}.pth'\n    matches = list(search_root.rglob(marker))\n    if matches:\n        return matches[0].parent\n    raise FileNotFoundError(\n        f'Could not find {marker} under {search_root}. '\n        f'Make sure the v07{version[-1]} notebook output is attached as an input dataset.'\n    )\n\nMODEL_A_DIR = find_model_dir('v07a')\nMODEL_B_DIR = find_model_dir('v07b')\nMODEL_F_DIR = find_model_dir('v07f')\nprint(f'v07a → {MODEL_A_DIR}')\nprint(f'v07b → {MODEL_B_DIR}')\nprint(f'v07f → {MODEL_F_DIR}')\n\n# Model configs for reconstruction\n# All 3 models output 1024-dim embeddings → clean averaging in ensemble\nMODEL_CONFIGS = {\n    'v07a': {\n        'model_name': 'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k',\n        'img_size': 448, 'batch_size': 4, 'grad_accum': 4,\n        'model_class': 'eva', 'input_dir': MODEL_A_DIR,\n        'norm_mean': [0.481, 0.457, 0.408], 'norm_std': [0.268, 0.261, 0.275],\n    },\n    'v07b': {\n        'model_name': 'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k',\n        'img_size': 448, 'batch_size': 4, 'grad_accum': 4,\n        'model_class': 'eva', 'input_dir': MODEL_B_DIR,\n        'norm_mean': [0.481, 0.457, 0.408], 'norm_std': [0.268, 0.261, 0.275],\n    },\n    'v07f': {\n        'model_name': 'vit_large_patch14_reg4_dinov2.lvd142m',\n        'img_size': 448, 'batch_size': 4, 'grad_accum': 4,\n        'model_class': 'dinov2', 'input_dir': MODEL_F_DIR,\n        # DINOv2 was trained with standard ImageNet normalization\n        'norm_mean': [0.485, 0.456, 0.406], 'norm_std': [0.229, 0.224, 0.225],\n    },\n}\n\n# ── Reproducibility ──────────────────────────────────────────────────────\n\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n\nseed_everything(SEED)\n\n# ── Dataset ──────────────────────────────────────────────────────────────\n\nclass JaguarDataset(Dataset):\n    \"\"\"Dataset that supports images from different directories (train + test).\"\"\"\n    def __init__(self, df, default_img_dir, transform=None, is_test=False):\n        self.df = df.reset_index(drop=True)\n        self.default_img_dir = Path(default_img_dir)\n        self.transform = transform\n        self.is_test = is_test\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_name = row['filename']\n        # Use per-row directory if available (for mixed train+pseudo data)\n        if 'dir' in row and pd.notna(row['dir']):\n            img_path = Path(row['dir']) / img_name\n        else:\n            img_path = self.default_img_dir / img_name\n        try:\n            img = Image.open(img_path).convert('RGB')\n        except Exception:\n            img = Image.new('RGB', (448, 448))\n        if self.transform:\n            img = self.transform(img)\n        if self.is_test:\n            return img, img_name\n        return img, torch.tensor(row['label'], dtype=torch.long)\n\n\ndef get_transforms(img_size, mode='train', norm_mean=None, norm_std=None):\n    \"\"\"Build transforms for a given image size and mode.\n    norm_mean/norm_std default to EVA jaguar stats; pass ImageNet stats for DINOv2.\"\"\"\n    mean = norm_mean if norm_mean is not None else NORM_MEAN\n    std = norm_std if norm_std is not None else NORM_STD\n    if mode == 'train':\n        return transforms.Compose([\n            transforms.Resize((img_size, img_size)),\n            transforms.RandomHorizontalFlip(),\n            transforms.RandomAffine(degrees=15, translate=(0.15, 0.15),\n                                    scale=(0.85, 1.15)),\n            transforms.ColorJitter(brightness=0.2, contrast=0.2),\n            transforms.ToTensor(),\n            transforms.Normalize(mean, std),\n            transforms.RandomErasing(p=0.25),\n        ])\n    else:\n        return transforms.Compose([\n            transforms.Resize((img_size, img_size)),\n            transforms.ToTensor(),\n            transforms.Normalize(mean, std),\n        ])\n\n\n# ── Model definitions ───────────────────────────────────────────────────\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.avg_pool2d(\n            x.clamp(min=self.eps).pow(self.p),\n            (x.size(-2), x.size(-1))\n        ).pow(1.0 / self.p)\n\n\nclass ArcFaceLayer(nn.Module):\n    def __init__(self, in_features, out_features, s=30.0, m=0.5):\n        super().__init__()\n        self.s = s\n        self.m = m\n        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))\n        nn.init.xavier_uniform_(self.weight)\n\n    def forward(self, input, label=None):\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n        if label is None:\n            return cosine\n        phi = cosine - self.m\n        one_hot = torch.zeros_like(cosine)\n        one_hot.scatter_(1, label.view(-1, 1), 1)\n        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n        return output * self.s\n\n\nclass EVAReIDModel(nn.Module):\n    def __init__(self, model_name, num_classes=31):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=True, num_classes=0)\n        self.feat_dim = self.backbone.num_features\n        self.gem = GeM()\n        self.bn = nn.BatchNorm1d(self.feat_dim)\n        self.head = ArcFaceLayer(self.feat_dim, num_classes, s=ARCFACE_S, m=ARCFACE_M)\n\n    def forward(self, x, label=None):\n        features = self.backbone.forward_features(x)\n        if features.dim() == 3:\n            B, N, C = features.shape\n            H = W = int(math.sqrt(N))\n            if H * W != N:\n                features = features[:, -H * W:, :]\n            features = features.permute(0, 2, 1).reshape(B, C, H, W)\n        emb = self.gem(features).flatten(1)\n        emb = self.bn(emb)\n        if label is not None:\n            return self.head(emb, label)\n        return emb\n\n\nclass DINOv2ReIDModel(nn.Module):\n    \"\"\"DINOv2-Large with register tokens.\n\n    Strips CLS (index 0) and 4 register tokens (indices 1-4) before\n    GeM pooling. At 448px with patch14: 32×32=1024 clean patch tokens.\"\"\"\n\n    NUM_REGISTER_TOKENS = 4\n\n    def __init__(self, model_name, num_classes=31):\n        super().__init__()\n        # img_size=448 MUST match the resolution used during v07f training.\n        # Without it, timm defaults to DINOv2's native 518px → 1369 pos_embed tokens,\n        # which is incompatible with the 448px checkpoint (1024 tokens).\n        self.backbone = timm.create_model(model_name, pretrained=True, num_classes=0, img_size=448)\n        self.feat_dim = self.backbone.num_features\n        self.gem = GeM()\n        self.bn = nn.BatchNorm1d(self.feat_dim)\n        self.head = ArcFaceLayer(self.feat_dim, num_classes, s=ARCFACE_S, m=ARCFACE_M)\n        self.patch_size = 14\n        self.grid_size = 448 // self.patch_size  # 32\n\n    def forward(self, x, label=None):\n        features = self.backbone.forward_features(x)\n        if features.dim() == 3:\n            B, total_tokens, C = features.shape\n            num_prefix = 1 + self.NUM_REGISTER_TOKENS\n            patch_tokens = features[:, num_prefix:, :]\n            H = W = self.grid_size\n            if patch_tokens.shape[1] != H * W:\n                patch_tokens = patch_tokens[:, -H * W:, :]\n            features = patch_tokens.permute(0, 2, 1).reshape(B, C, H, W)\n        emb = self.gem(features).flatten(1)\n        emb = self.bn(emb)\n        if label is not None:\n            return self.head(emb, label)\n        return emb\n\n\ndef build_model(model_class, model_name, num_classes=31):\n    \"\"\"Factory: build model by class type.\"\"\"\n    if model_class == 'eva':\n        return EVAReIDModel(model_name, num_classes)\n    elif model_class == 'dinov2':\n        return DINOv2ReIDModel(model_name, num_classes)\n    else:\n        raise ValueError(f'Unknown model class: {model_class}')\n\n\n# ── EMA ──────────────────────────────────────────────────────────────────\n\nclass EMAModel:\n    def __init__(self, model, decay=0.999):\n        self.decay = decay\n        self.shadow = deepcopy(model)\n        self.shadow.eval()\n        for p in self.shadow.parameters():\n            p.requires_grad_(False)\n\n    @torch.no_grad()\n    def update(self, model):\n        for ema_p, model_p in zip(self.shadow.parameters(), model.parameters()):\n            ema_p.data.mul_(self.decay).add_(model_p.data, alpha=1 - self.decay)\n        for ema_b, model_b in zip(self.shadow.buffers(), model.buffers()):\n            ema_b.data.copy_(model_b.data)\n\n    def state_dict(self):\n        return self.shadow.state_dict()\n\n\n# ── Pseudo-label generation ─────────────────────────────────────────────\n\ndef compute_class_centroids(train_emb, train_labels, num_classes=31):\n    \"\"\"Compute L2-normalised centroid for each class from training embeddings.\"\"\"\n    centroids = np.zeros((num_classes, train_emb.shape[1]), dtype=np.float32)\n    for c in range(num_classes):\n        mask = train_labels == c\n        if mask.sum() > 0:\n            centroids[c] = train_emb[mask].mean(axis=0)\n    # L2 normalise\n    norms = np.linalg.norm(centroids, axis=1, keepdims=True)\n    norms = np.maximum(norms, 1e-8)\n    centroids = centroids / norms\n    return centroids\n\n\ndef generate_pseudo_labels(ensemble_test_emb, ensemble_train_emb, train_labels,\n                           test_filenames, threshold=0.90, max_add=500):\n    \"\"\"Generate pseudo-labels for test images using ensemble embeddings.\n\n    Computes class centroids from training embeddings, then assigns each test\n    image the class whose centroid has highest cosine similarity, if above threshold.\n    \"\"\"\n    centroids = compute_class_centroids(ensemble_train_emb, train_labels)\n\n    # Cosine similarity: test_emb @ centroids.T → (N_test, N_classes)\n    sims = ensemble_test_emb @ centroids.T\n\n    max_sims = sims.max(axis=1)\n    pred_classes = sims.argmax(axis=1)\n\n    pl_data = []\n    for i in range(len(max_sims)):\n        if max_sims[i] > threshold and len(pl_data) < max_add:\n            pl_data.append({\n                'filename': test_filenames[i],\n                'label': int(pred_classes[i]),\n                'confidence': float(max_sims[i]),\n                'dir': str(TEST_DIR),\n            })\n\n    print(f'  Pseudo-labels: {len(pl_data)}/{len(test_filenames)} test images '\n          f'above threshold {threshold}')\n    if pl_data:\n        confs = [d['confidence'] for d in pl_data]\n        print(f'  Confidence range: {min(confs):.4f} – {max(confs):.4f}')\n\n    return pd.DataFrame(pl_data)\n\n\n# ── Training ─────────────────────────────────────────────────────────────\n\ndef train_epoch(model, loader, optimizer, criterion, scaler, ema, grad_accum):\n    model.train()\n    total_loss = 0.0\n    optimizer.zero_grad()\n    for i, (imgs, labels) in enumerate(tqdm(loader, leave=False, desc='PL Train')):\n        imgs = imgs.to(DEVICE)\n        labels = labels.to(DEVICE)\n        with torch.amp.autocast(DEVICE_TYPE):\n            logits = model(imgs, labels)\n            loss = criterion(logits, labels)\n            loss = loss / grad_accum\n        scaler.scale(loss).backward()\n        if (i + 1) % grad_accum == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n            ema.update(model)\n        total_loss += loss.item() * grad_accum\n    if (i + 1) % grad_accum != 0:\n        scaler.step(optimizer)\n        scaler.update()\n        optimizer.zero_grad()\n        ema.update(model)\n    return total_loss / len(loader)\n\n\n# ── Inference ────────────────────────────────────────────────────────────\n\n@torch.no_grad()\ndef extract_features(model, loader, use_tta=True):\n    model.eval()\n    all_feats, all_names = [], []\n    for imgs, fnames in tqdm(loader, desc='Extracting'):\n        imgs = imgs.to(DEVICE)\n        with torch.amp.autocast(DEVICE_TYPE):\n            f1 = model(imgs)\n            if use_tta:\n                f2 = model(torch.flip(imgs, [3]))\n                f1 = (f1 + f2) / 2\n        all_feats.append(F.normalize(f1, dim=1).float().cpu())\n        all_names.extend(fnames)\n    return torch.cat(all_feats, dim=0).numpy(), all_names\n\n\n# ── Post-processing (for standalone submission) ──────────────────────────\n\ndef query_expansion(emb, top_k=3):\n    sims = emb @ emb.T\n    indices = np.argsort(-sims, axis=1)[:, :top_k]\n    new_emb = np.zeros_like(emb)\n    for i in range(len(emb)):\n        new_emb[i] = np.mean(emb[indices[i]], axis=0)\n    return new_emb / np.linalg.norm(new_emb, axis=1, keepdims=True)\n\n\ndef k_reciprocal_rerank(prob, k1=15, k2=6, lambda_value=0.4):\n    q_g_dist = 1 - prob\n    original_dist = q_g_dist.copy()\n    initial_rank = np.argsort(original_dist, axis=1)\n    nn_k1 = []\n    for i in range(prob.shape[0]):\n        forward_k1 = initial_rank[i, :k1 + 1]\n        backward_k1 = initial_rank[forward_k1, :k1 + 1]\n        fi = np.where(backward_k1 == i)[0]\n        nn_k1.append(forward_k1[fi])\n    jaccard_dist = np.zeros_like(original_dist)\n    for i in range(prob.shape[0]):\n        ind_non_zero = np.where(original_dist[i, :] < 0.6)[0]\n        ind_images = [\n            inv for inv in ind_non_zero\n            if len(np.intersect1d(nn_k1[i], nn_k1[inv])) > 0\n        ]\n        for j in ind_images:\n            intersection = len(np.intersect1d(nn_k1[i], nn_k1[j]))\n            union = len(np.union1d(nn_k1[i], nn_k1[j]))\n            jaccard_dist[i, j] = 1 - intersection / union\n    return 1 - (jaccard_dist * lambda_value + original_dist * (1 - lambda_value))\n\n\n# ══════════════════════════════════════════════════════════════════════════\n# MAIN EXECUTION\n# ══════════════════════════════════════════════════════════════════════════\n\nif __name__ == '__main__':\n    t0 = time.time()\n\n    train_df = pd.read_csv(KAGGLE_INPUT / 'train.csv')\n    test_df = pd.read_csv(KAGGLE_INPUT / 'test.csv')\n\n    # Encode labels\n    unique_ids = sorted(train_df['ground_truth'].unique())\n    label_map = {name: i for i, name in enumerate(unique_ids)}\n    idx_to_label = {i: name for name, i in label_map.items()}\n    train_df['label'] = train_df['ground_truth'].map(label_map)\n    train_df['dir'] = str(TRAIN_DIR)\n\n    print(f'Train: {len(train_df)} images, {len(unique_ids)} classes')\n    print(f'Test pairs: {len(test_df)}')\n\n    # ── Step 1: Load pre-computed embeddings from base models ────────────\n    print('\\n' + '='*60)\n    print('STEP 1: Loading base model embeddings')\n    print('='*60)\n\n    all_test_embs = {}\n    all_train_embs = {}\n    test_filenames = None\n\n    for version, cfg in MODEL_CONFIGS.items():\n        input_dir = cfg['input_dir']\n        print(f'\\n  Loading {version} from {input_dir}...')\n\n        test_emb = np.load(input_dir / f'embeddings_{version}_test.npy')\n        train_emb = np.load(input_dir / f'embeddings_{version}_train.npy')\n        with open(input_dir / f'filenames_{version}_test.json') as f:\n            test_names = json.load(f)\n        with open(input_dir / f'filenames_{version}_train.json') as f:\n            train_names = json.load(f)\n\n        all_test_embs[version] = test_emb\n        all_train_embs[version] = train_emb\n        print(f'    Test:  {test_emb.shape}')\n        print(f'    Train: {train_emb.shape}')\n\n        if test_filenames is None:\n            test_filenames = test_names\n            train_filenames = train_names\n\n    # ── Step 2: Ensemble embeddings for pseudo-labeling ──────────────────\n    print('\\n' + '='*60)\n    print('STEP 2: Ensembling embeddings for pseudo-label generation')\n    print('='*60)\n\n    # All 3 models share 1024-dim space → average all three for strongest signal\n    def l2_norm(e):\n        return e / np.maximum(np.linalg.norm(e, axis=1, keepdims=True), 1e-8)\n\n    ens_test = l2_norm(\n        all_test_embs['v07a'] + all_test_embs['v07b'] + all_test_embs['v07f'])\n    ens_train = l2_norm(\n        all_train_embs['v07a'] + all_train_embs['v07b'] + all_train_embs['v07f'])\n    print(f'  3-model average: test {ens_test.shape}, train {ens_train.shape}')\n\n    # Per-model predictions for cross-validation\n    per_model_test = {v: all_test_embs[v] for v in MODEL_CONFIGS}\n    per_model_train = {v: all_train_embs[v] for v in MODEL_CONFIGS}\n\n    # Get train labels in embedding order\n    train_label_map = {fn: label_map[train_df.loc[train_df['filename'] == fn,\n                       'ground_truth'].values[0]]\n                       for fn in train_filenames}\n    train_labels = np.array([train_label_map[fn] for fn in train_filenames])\n\n    # ── Step 3: Generate pseudo-labels ───────────────────────────────────\n    print('\\n' + '='*60)\n    print('STEP 3: Generating pseudo-labels (3-way agreement)')\n    print('='*60)\n\n    # Generate per-model pseudo-labels for cross-validation\n    per_model_pls = {}\n    for version in MODEL_CONFIGS:\n        pl_df_v = generate_pseudo_labels(\n            per_model_test[version], per_model_train[version],\n            train_labels, test_filenames,\n            threshold=PL_THRESHOLD, max_add=PL_MAX_ADD)\n        per_model_pls[version] = dict(zip(pl_df_v['filename'], pl_df_v['label'])) \\\n            if len(pl_df_v) > 0 else {}\n        print(f'  {version}: {len(per_model_pls[version])} above threshold')\n\n    # Also run on 3-model average for highest-confidence candidates\n    pl_df_ens = generate_pseudo_labels(\n        ens_test, ens_train, train_labels, test_filenames,\n        threshold=PL_THRESHOLD, max_add=PL_MAX_ADD)\n    ens_pls = dict(zip(pl_df_ens['filename'], pl_df_ens['label'])) \\\n        if len(pl_df_ens) > 0 else {}\n\n    # Keep only pseudo-labels where ALL 3 individual models agree\n    pls_a, pls_b, pls_f = (per_model_pls['v07a'],\n                            per_model_pls['v07b'],\n                            per_model_pls['v07f'])\n    agreed = []\n    disagreed = 0\n    for fn in ens_pls:\n        label_a = pls_a.get(fn)\n        label_b = pls_b.get(fn)\n        label_f = pls_f.get(fn)\n        if label_a is not None and label_b is not None and label_f is not None:\n            if label_a == label_b == label_f:\n                row = pl_df_ens[pl_df_ens['filename'] == fn].iloc[0].to_dict()\n                agreed.append(row)\n            else:\n                disagreed += 1\n\n    pl_df = pd.DataFrame(agreed) if agreed else pd.DataFrame()\n    print(f'\\n  3-way agreement: {len(agreed)} agreed, {disagreed} disagreed')\n\n    # Save pseudo-labels for analysis\n    if len(pl_df) > 0:\n        pl_df['identity'] = pl_df['label'].map(idx_to_label)\n        pl_df.to_csv(OUT_DIR / 'pseudo_labels.csv', index=False)\n        print(f'  Saved pseudo_labels.csv ({len(pl_df)} entries)')\n\n        # Distribution\n        print(f'\\n  Pseudo-label distribution:')\n        for label_idx in sorted(pl_df['label'].unique()):\n            count = (pl_df['label'] == label_idx).sum()\n            name = idx_to_label[label_idx]\n            print(f'    {name}: {count}')\n\n    # ── Step 4: Fine-tune each model on train + pseudo-labels ────────────\n    print('\\n' + '='*60)\n    print('STEP 4: Fine-tuning models with pseudo-labels')\n    print('='*60)\n\n    if len(pl_df) == 0:\n        print('  No pseudo-labels generated! Skipping fine-tuning.')\n        print('  Copying base model embeddings as-is.')\n        for version in MODEL_CONFIGS:\n            cfg = MODEL_CONFIGS[version]\n            input_dir = cfg['input_dir']\n            # Just copy embeddings with _pl suffix\n            for suffix in ['test', 'train']:\n                src = np.load(input_dir / f'embeddings_{version}_{suffix}.npy')\n                np.save(OUT_DIR / f'embeddings_{version}_pl_{suffix}.npy', src)\n            for suffix in ['test', 'train']:\n                with open(input_dir / f'filenames_{version}_{suffix}.json') as f:\n                    data = json.load(f)\n                with open(OUT_DIR / f'filenames_{version}_pl_{suffix}.json', 'w') as fw:\n                    json.dump(data, fw)\n    else:\n        # Prepare combined dataset\n        combined_base = train_df[['filename', 'label', 'dir']].copy()\n        pl_combined = pl_df[['filename', 'label', 'dir']].copy()\n        combined_df = pd.concat([combined_base, pl_combined], ignore_index=True)\n        print(f'  Combined dataset: {len(combined_df)} images '\n              f'({len(combined_base)} train + {len(pl_combined)} pseudo)')\n\n        unique_test_fnames = sorted(\n            set(test_df['query_image']) | set(test_df['gallery_image']))\n\n        for version, cfg in MODEL_CONFIGS.items():\n            print(f'\\n  {\"─\"*50}')\n            print(f'  Fine-tuning {version}: {cfg[\"model_name\"]}')\n            print(f'  {\"─\"*50}')\n\n            img_size = cfg['img_size']\n            batch_size = cfg['batch_size']\n            grad_accum = cfg['grad_accum']\n\n            # Build model and load base weights\n            model = build_model(cfg['model_class'], cfg['model_name']).to(DEVICE)\n            if hasattr(model.backbone, 'set_grad_checkpointing'):\n                model.backbone.set_grad_checkpointing(True)\n\n            base_weights = torch.load(\n                cfg['input_dir'] / f'model_{version}.pth',\n                map_location=DEVICE)\n            model.load_state_dict(base_weights)\n            print(f'    Loaded base weights from {cfg[\"input_dir\"]}')\n\n            # Create dataloaders (use per-model normalization)\n            norm_mean = cfg.get('norm_mean')\n            norm_std = cfg.get('norm_std')\n            train_transform = get_transforms(img_size, mode='train', norm_mean=norm_mean, norm_std=norm_std)\n            test_transform = get_transforms(img_size, mode='test', norm_mean=norm_mean, norm_std=norm_std)\n\n            train_dataset = JaguarDataset(combined_df, TRAIN_DIR, train_transform)\n            train_loader = DataLoader(\n                train_dataset, batch_size=batch_size, shuffle=True,\n                num_workers=2, pin_memory=True, drop_last=True)\n\n            test_dataset = JaguarDataset(\n                pd.DataFrame({'filename': unique_test_fnames}),\n                TEST_DIR, test_transform, is_test=True)\n            test_loader = DataLoader(\n                test_dataset, batch_size=batch_size, shuffle=False,\n                num_workers=2, pin_memory=True)\n\n            # Optimizer + EMA\n            optimizer = torch.optim.AdamW(\n                model.parameters(), lr=PL_LR, weight_decay=WEIGHT_DECAY)\n            scaler = torch.amp.GradScaler(DEVICE_TYPE)\n            ema = EMAModel(model, decay=EMA_DECAY)\n            criterion = nn.CrossEntropyLoss()\n\n            # Fine-tune\n            for epoch in range(PL_EPOCHS):\n                loss = train_epoch(\n                    model, train_loader, optimizer, criterion,\n                    scaler, ema, grad_accum)\n                print(f'    Epoch {epoch+1}/{PL_EPOCHS} | Loss: {loss:.4f}')\n\n            # Save PL model weights\n            model_path = OUT_DIR / f'model_{version}_pl.pth'\n            torch.save(ema.state_dict(), model_path)\n            print(f'    Saved → {model_path}')\n\n            # Extract updated embeddings using EMA model\n            ema_model = ema.shadow\n\n            test_emb, test_names = extract_features(\n                ema_model, test_loader, use_tta=USE_TTA)\n\n            # Train embeddings\n            train_emb_dataset = JaguarDataset(\n                pd.DataFrame({'filename': train_df['filename'].tolist()}),\n                TRAIN_DIR, test_transform, is_test=True)\n            train_emb_loader = DataLoader(\n                train_emb_dataset, batch_size=batch_size, shuffle=False,\n                num_workers=2, pin_memory=True)\n            train_emb, train_names = extract_features(\n                ema_model, train_emb_loader, use_tta=USE_TTA)\n\n            # Save\n            np.save(OUT_DIR / f'embeddings_{version}_pl_test.npy', test_emb)\n            np.save(OUT_DIR / f'embeddings_{version}_pl_train.npy', train_emb)\n            with open(OUT_DIR / f'filenames_{version}_pl_test.json', 'w') as f:\n                json.dump(test_names, f)\n            with open(OUT_DIR / f'filenames_{version}_pl_train.json', 'w') as f:\n                json.dump(train_names, f)\n            print(f'    Embeddings: test {test_emb.shape}, train {train_emb.shape}')\n\n            # Free GPU memory before next model\n            del model, ema, optimizer, scaler\n            torch.cuda.empty_cache()\n\n    # ── Step 5: Generate standalone submission from best PL model ────────\n    print('\\n' + '='*60)\n    print('STEP 5: Generating standalone submission (EVA-A PL)')\n    print('='*60)\n\n    # Use the v07a PL model as the standalone submission\n    test_emb = np.load(OUT_DIR / 'embeddings_v07a_pl_test.npy')\n    with open(OUT_DIR / 'filenames_v07a_pl_test.json') as f:\n        test_names = json.load(f)\n    img_map = {n: i for i, n in enumerate(test_names)}\n\n    # QE + rerank\n    emb = query_expansion(test_emb, top_k=3)\n    sim_matrix = emb @ emb.T\n    sim_matrix = k_reciprocal_rerank(sim_matrix)\n\n    preds = []\n    for _, row in tqdm(test_df.iterrows(), total=len(test_df), desc='Mapping'):\n        s = sim_matrix[img_map[row['query_image']], img_map[row['gallery_image']]]\n        preds.append(max(0.0, min(1.0, float(s))))\n\n    sub = pd.DataFrame({'row_id': test_df['row_id'], 'similarity': preds})\n    sub.to_csv(OUT_DIR / f'submission_{VERSION}.csv', index=False)\n    sub.to_csv(Path('/kaggle/working/submission.csv'), index=False)\n    print(f'  Score range: {np.min(preds):.4f} – {np.max(preds):.4f}')\n    print(f'  Score mean:  {np.mean(preds):.4f}')\n\n    total_time = time.time() - t0\n    print(f'\\n{\"=\"*60}')\n    print(f'{VERSION} complete! Total: {total_time/60:.1f} min')\n    print(f'{\"=\"*60}')","metadata":{"_uuid":"65863aed-edb3-494d-9ba4-2964aa946076","_cell_guid":"ab5a940c-e500-4822-bb37-c22119d48f80","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null}]}