{"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}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nv07a: EVA-02 Large @ 448px — Model A (seed=42)\n================================================\nPhase 1 of the v07 ensemble pipeline targeting 0.97+ LB.\n\nArchitecture:\n  - EVA-02 Large (300M params, patch14 @ 448px = 1024 tokens)\n  - GeM pooling (learnable p=3) → BatchNorm → ArcFace (s=30, m=0.5)\n  - Gradient checkpointing to fit 448px into 16GB VRAM\n  - EMA (decay=0.999) for smoother inference weights\n\nTraining:\n  - Full train set (1895 images, 31 classes), no validation split\n  - AdamW lr=2e-5, CosineAnnealingLR, 10 epochs\n  - Batch=4, grad_accum=4 → effective batch=16\n  - Augmentation: HFlip, RandomAffine(15°, 0.15, 0.85-1.15), ColorJitter, RandomErasing\n\nOutputs saved to /kaggle/working/output/:\n  - model_v07a.pth             (EMA weights for pseudo-labeling notebook)\n  - embeddings_v07a_train.npy  (1895, 1024) — train embeddings\n  - embeddings_v07a_test.npy   (371, 1024)  — test embeddings\n  - filenames_v07a_train.json  (ordered filenames matching embedding rows)\n  - filenames_v07a_test.json\n  - submission_v07a.csv        (standalone submission with QE + rerank)\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 = 'v07a'\nSEED = 42\n\n# Model\nMODEL_NAME = 'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k'\nIMG_SIZE = 448\nNUM_CLASSES = 31\n\n# Training\nNUM_EPOCHS = 20\nBATCH_SIZE = 4\nGRAD_ACCUM = 4       # effective batch = 16\nLR = 2e-5\nWEIGHT_DECAY = 1e-3\nEMA_DECAY = 0.999\n\n# ArcFace\nARCFACE_S = 30.0\nARCFACE_M = 0.5\n\n# Post-processing\nUSE_TTA = True        # horizontal flip TTA\nUSE_QE = True         # query expansion\nQE_TOP_K = 3\nUSE_RERANK = True     # k-reciprocal reranking\nRERANK_K1 = 15        # tuned for small 371-image gallery (default 20)\nRERANK_K2 = 6\nRERANK_LAMBDA = 0.4   # trust Jaccard more with small gallery (default 0.3)\n\n# Normalisation (from reference EVA notebook — computed on jaguar dataset)\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# ── 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    def __init__(self, df, img_dir, transform=None, is_test=False):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = Path(img_dir)\n        self.transform = transform\n        self.is_test = is_test\n        if not is_test:\n            unique_ids = sorted(df['ground_truth'].unique())\n            self.label_map = {name: i for i, name in enumerate(unique_ids)}\n            self.df['label'] = self.df['ground_truth'].map(self.label_map)\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        img_path = self.img_dir / img_name\n        try:\n            img = Image.open(img_path).convert('RGB')\n        except Exception:\n            img = Image.new('RGB', (IMG_SIZE, IMG_SIZE))\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# ── Transforms ───────────────────────────────────────────────────────────\n\ntrain_transform = transforms.Compose([\n    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n    transforms.RandomHorizontalFlip(),\n    transforms.RandomAffine(degrees=15, translate=(0.15, 0.15), scale=(0.85, 1.15)),\n    transforms.ColorJitter(brightness=0.2, contrast=0.2),\n    transforms.ToTensor(),\n    transforms.Normalize(NORM_MEAN, NORM_STD),\n    transforms.RandomErasing(p=0.25),\n])\n\ntest_transform = transforms.Compose([\n    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n    transforms.ToTensor(),\n    transforms.Normalize(NORM_MEAN, NORM_STD),\n])\n\n# ── Model ────────────────────────────────────────────────────────────────\n\nclass GeM(nn.Module):\n    \"\"\"Generalised Mean Pooling. p=1 → avg pool, p→∞ → max pool.\n    Learnable p starting at 3 upweights high-activation spatial regions\n    (distinctive spot clusters) while still aggregating context.\"\"\"\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    \"\"\"ArcFace angular margin head.\n    During training: adds margin m to the angle between embedding and true-class\n    weight, then scales by s before CrossEntropyLoss.\n    During inference (label=None): returns raw cosine logits (unused — we extract\n    embeddings from the layer before this).\"\"\"\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    \"\"\"EVA-02 Large backbone with GeM pooling and ArcFace head.\n\n    forward_features() returns (B, N, C) tokens for ViTs. We reshape to\n    (B, C, H, W) spatial grid so GeM can pool spatially. For 448px with\n    patch14: N = 32×32 = 1024 tokens, C = 1024 channels.\n\n    During training: returns ArcFace logits (feed to CrossEntropyLoss).\n    During inference (label=None): returns L2-normalised embeddings.\"\"\"\n\n    def __init__(self):\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        # ViT: (B, N, C) → (B, C, H, W) for spatial GeM pooling\n        if features.dim() == 3:\n            B, N, C = features.shape\n            H = W = int(math.sqrt(N))\n            if H * W != N:\n                # Some ViTs prepend a CLS token; take last H*W tokens\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\n# ── EMA ──────────────────────────────────────────────────────────────────\n\nclass EMAModel:\n    \"\"\"Exponential Moving Average of model weights.\n    Shadow parameters are updated each step: ema = decay * ema + (1-decay) * model.\n    Produces a smoother model that generalizes better than any single checkpoint.\"\"\"\n\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        # Also update buffers (BatchNorm running mean/var)\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# ── Training ─────────────────────────────────────────────────────────────\n\ndef train_epoch(model, loader, optimizer, criterion, scaler, ema):\n    model.train()\n    total_loss = 0.0\n    optimizer.zero_grad()\n    for i, (imgs, labels) in enumerate(tqdm(loader, leave=False, desc='Train')):\n        imgs = imgs.to(DEVICE)\n        labels = labels.to(DEVICE)\n\n        with torch.amp.autocast(DEVICE_TYPE):\n            logits = model(imgs, labels)\n            loss = criterion(logits, labels)\n            loss = loss / GRAD_ACCUM\n\n        scaler.scale(loss).backward()\n\n        if (i + 1) % GRAD_ACCUM == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n            ema.update(model)\n\n        total_loss += loss.item() * GRAD_ACCUM\n\n    # Handle leftover steps (if len(loader) not divisible by 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\n    return total_loss / len(loader)\n\n\n# ── Inference ────────────────────────────────────────────────────────────\n\n@torch.no_grad()\ndef extract_features(model, loader, use_tta=True):\n    \"\"\"Extract L2-normalised embeddings. TTA = horizontal flip averaging.\"\"\"\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@torch.no_grad()\ndef extract_train_features(model, filenames, img_dir, transform, batch_size=4):\n    \"\"\"Extract embeddings for training images (returns embeddings in filename order).\"\"\"\n    model.eval()\n    df = pd.DataFrame({'filename': filenames})\n    dataset = JaguarDataset(df, img_dir, transform, is_test=True)\n    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False,\n                        num_workers=2, pin_memory=True)\n    return extract_features(model, loader, use_tta=USE_TTA)\n\n\n# ── Post-processing ─────────────────────────────────────────────────────\n\ndef query_expansion(emb, top_k=3):\n    \"\"\"Replace each embedding with the mean of its top-k neighbours.\n    This 'expands' the query to be more representative of its local cluster,\n    improving recall for hard positives.\"\"\"\n    print(f'  Applying Query Expansion (top_k={top_k})...')\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=20, k2=6, lambda_value=0.3):\n    \"\"\"K-reciprocal reranking (Zhong et al., CVPR 2017).\n    Uses Jaccard distance between k-reciprocal nearest neighbour sets\n    to refine the initial cosine similarity ranking. The key idea:\n    if A is in B's top-k AND B is in A's top-k, they're likely the same identity.\"\"\"\n    print(f'  Applying k-reciprocal reranking (k1={k1}, k2={k2}, λ={lambda_value})...')\n    q_g_dist = 1 - prob\n    original_dist = q_g_dist.copy()\n    initial_rank = np.argsort(original_dist, axis=1)\n\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\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\n    return 1 - (jaccard_dist * lambda_value + original_dist * (1 - lambda_value))\n\n\n# ── Submission ───────────────────────────────────────────────────────────\n\ndef generate_submission(test_df, emb, img_map, output_path):\n    \"\"\"Generate submission CSV from embedding similarity matrix.\"\"\"\n    if USE_QE:\n        emb = query_expansion(emb, top_k=QE_TOP_K)\n\n    sim_matrix = emb @ emb.T\n\n    if USE_RERANK:\n        sim_matrix = k_reciprocal_rerank(\n            sim_matrix, k1=RERANK_K1, k2=RERANK_K2, lambda_value=RERANK_LAMBDA)\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(output_path, index=False)\n    print(f'  Submission saved → {output_path}')\n    print(f'  Score range: {np.min(preds):.4f} – {np.max(preds):.4f}')\n    print(f'  Score mean:  {np.mean(preds):.4f}')\n    return sub\n\n\n# ══════════════════════════════════════════════════════════════════════════\n# MAIN EXECUTION\n# ══════════════════════════════════════════════════════════════════════════\n\nif __name__ == '__main__':\n    t0 = time.time()\n\n    # ── Load data ────────────────────────────────────────────────────────\n    train_df = pd.read_csv(KAGGLE_INPUT / 'train.csv')\n    test_df = pd.read_csv(KAGGLE_INPUT / 'test.csv')\n    print(f'Train: {len(train_df)} images, {train_df[\"ground_truth\"].nunique()} classes')\n    print(f'Test pairs: {len(test_df)}')\n\n    # ── Create dataloaders ───────────────────────────────────────────────\n    train_dataset = JaguarDataset(train_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    unique_test = sorted(set(test_df['query_image']) | set(test_df['gallery_image']))\n    test_dataset = JaguarDataset(\n        pd.DataFrame({'filename': unique_test}), 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    # ── Build model ──────────────────────────────────────────────────────\n    model = EVAReIDModel().to(DEVICE)\n\n    # Enable gradient checkpointing to save ~40% VRAM\n    if hasattr(model.backbone, 'set_grad_checkpointing'):\n        model.backbone.set_grad_checkpointing(True)\n        print('Gradient checkpointing: ENABLED')\n\n    n_params = sum(p.numel() for p in model.parameters())\n    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n    print(f'Model: {MODEL_NAME}')\n    print(f'Parameters: {n_trainable:,} trainable / {n_params:,} total')\n    print(f'Embedding dim: {model.feat_dim}')\n\n    # ── Optimizer, scheduler, scaler, EMA ────────────────────────────────\n    optimizer = torch.optim.AdamW(\n        model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer, T_max=NUM_EPOCHS)\n    scaler = torch.amp.GradScaler(DEVICE_TYPE)\n    ema = EMAModel(model, decay=EMA_DECAY)\n    criterion = nn.CrossEntropyLoss()\n\n    # ── Train ────────────────────────────────────────────────────────────\n    print(f'\\n{\"=\"*60}')\n    print(f'Training {VERSION}: {MODEL_NAME} @ {IMG_SIZE}px')\n    print(f'Epochs: {NUM_EPOCHS} | Batch: {BATCH_SIZE}×{GRAD_ACCUM} | LR: {LR}')\n    print(f'{\"=\"*60}\\n')\n\n    for epoch in range(NUM_EPOCHS):\n        epoch_t0 = time.time()\n        loss = train_epoch(model, train_loader, optimizer, criterion, scaler, ema)\n        scheduler.step()\n        elapsed = time.time() - epoch_t0\n        lr_now = scheduler.get_last_lr()[0]\n        print(f'Epoch {epoch+1:2d}/{NUM_EPOCHS} | '\n              f'Loss: {loss:.4f} | LR: {lr_now:.2e} | '\n              f'Time: {elapsed:.0f}s')\n\n    total_train = time.time() - t0\n    print(f'\\nTraining complete in {total_train/60:.1f} min')\n\n    # ── Save model weights (EMA) ─────────────────────────────────────────\n    model_path = OUT_DIR / f'model_{VERSION}.pth'\n    torch.save(ema.state_dict(), model_path)\n    print(f'EMA weights saved → {model_path}')\n\n    # ── Extract embeddings using EMA model ───────────────────────────────\n    ema_model = ema.shadow\n\n    # Test embeddings\n    print('\\nExtracting test embeddings (EMA model)...')\n    test_emb, test_names = extract_features(ema_model, test_loader, use_tta=USE_TTA)\n    print(f'  Test embeddings: {test_emb.shape}')\n\n    # Train embeddings (for pseudo-labeling notebook)\n    print('Extracting train embeddings (EMA model)...')\n    train_emb, train_names = extract_train_features(\n        ema_model, train_df['filename'].tolist(), TRAIN_DIR, test_transform,\n        batch_size=BATCH_SIZE)\n    print(f'  Train embeddings: {train_emb.shape}')\n\n    # ── Save embeddings and filenames ────────────────────────────────────\n    np.save(OUT_DIR / f'embeddings_{VERSION}_test.npy', test_emb)\n    np.save(OUT_DIR / f'embeddings_{VERSION}_train.npy', train_emb)\n    with open(OUT_DIR / f'filenames_{VERSION}_test.json', 'w') as f:\n        json.dump(test_names, f)\n    with open(OUT_DIR / f'filenames_{VERSION}_train.json', 'w') as f:\n        json.dump(train_names, f)\n    print(f'Embeddings + filenames saved to {OUT_DIR}')\n\n    # ── Generate standalone submission ───────────────────────────────────\n    print('\\nGenerating standalone submission...')\n    img_map = {n: i for i, n in enumerate(test_names)}\n    generate_submission(\n        test_df, test_emb, img_map,\n        output_path=OUT_DIR / f'submission_{VERSION}.csv')\n\n    # Also save to /kaggle/working/ for easy download\n    generate_submission(\n        test_df, test_emb, img_map,\n        output_path=Path('/kaggle/working/submission.csv'))\n\n    total_time = time.time() - t0\n    print(f'\\n{\"=\"*60}')\n    print(f'{VERSION} complete! Total time: {total_time/60:.1f} min')\n    print(f'{\"=\"*60}')","metadata":{"_uuid":"15cbf2dd-7a8d-4859-a95b-7ca7b43e4d13","_cell_guid":"8ea99e66-94d3-4df4-a4c3-018d3ec1de5a","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null}]}