{"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}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nv07b: EVA-02 Large @ 448px — Model B (seed=123, different augmentation)\n=======================================================================\nPhase 1 of the v07 ensemble pipeline. Same architecture as v07a but with:\n  - seed=123 (different weight init, data order)\n  - Milder affine (10°, 0.1 translate, 0.9-1.1 scale)\n  - Stronger color jitter (0.3, 0.3)\n  - Added RandomAdjustSharpness (p=0.3)\nThese differences decorrelate error patterns for ensemble diversity.\n\nOutputs saved to /kaggle/working/output/:\n  - model_v07b.pth, embeddings_v07b_{train,test}.npy, filenames_v07b_{train,test}.json\n  - submission_v07b.csv (standalone)\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 = 'v07b'\nSEED = 123\n\nMODEL_NAME = 'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k'\nIMG_SIZE = 448\nNUM_CLASSES = 31\n\nNUM_EPOCHS = 20\nBATCH_SIZE = 4\nGRAD_ACCUM = 4\nLR = 2e-5\nWEIGHT_DECAY = 1e-3\nEMA_DECAY = 0.999\n\nARCFACE_S = 30.0\nARCFACE_M = 0.5\n\nUSE_TTA = True\nUSE_QE = True\nQE_TOP_K = 3\nUSE_RERANK = True\nRERANK_K1 = 15\nRERANK_K2 = 6\nRERANK_LAMBDA = 0.4\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# ── 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 (Model B: milder affine, stronger color, + sharpness) ────\n\ntrain_transform = transforms.Compose([\n    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n    transforms.RandomHorizontalFlip(),\n    transforms.RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.9, 1.1)),\n    transforms.ColorJitter(brightness=0.3, contrast=0.3),\n    transforms.RandomAdjustSharpness(sharpness_factor=2, p=0.3),\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    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):\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\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# ── 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        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@torch.no_grad()\ndef extract_train_features(model, filenames, img_dir, transform, batch_size=4):\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    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    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    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\ndef generate_submission(test_df, emb, img_map, output_path):\n    if USE_QE:\n        emb = query_expansion(emb, top_k=QE_TOP_K)\n    sim_matrix = emb @ emb.T\n    if USE_RERANK:\n        sim_matrix = k_reciprocal_rerank(\n            sim_matrix, k1=RERANK_K1, k2=RERANK_K2, lambda_value=RERANK_LAMBDA)\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    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\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    print(f'Train: {len(train_df)} images, {train_df[\"ground_truth\"].nunique()} classes')\n\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    model = EVAReIDModel().to(DEVICE)\n    if hasattr(model.backbone, 'set_grad_checkpointing'):\n        model.backbone.set_grad_checkpointing(True)\n        print('Gradient checkpointing: ENABLED')\n\n    print(f'Model: {MODEL_NAME} | Seed: {SEED}')\n    print(f'Params: {sum(p.numel() for p in model.parameters()):,}')\n\n    optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(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    print(f'\\n{\"=\"*60}')\n    print(f'Training {VERSION}: {MODEL_NAME} @ {IMG_SIZE}px (seed={SEED})')\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        print(f'Epoch {epoch+1:2d}/{NUM_EPOCHS} | '\n              f'Loss: {loss:.4f} | LR: {scheduler.get_last_lr()[0]:.2e} | '\n              f'Time: {time.time()-epoch_t0:.0f}s')\n\n    print(f'\\nTraining complete in {(time.time()-t0)/60:.1f} min')\n\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    ema_model = ema.shadow\n    print('\\nExtracting test embeddings...')\n    test_emb, test_names = extract_features(ema_model, test_loader, use_tta=USE_TTA)\n    print(f'  Test: {test_emb.shape}')\n\n    print('Extracting train embeddings...')\n    train_emb, train_names = extract_train_features(\n        ema_model, train_df['filename'].tolist(), TRAIN_DIR, test_transform, BATCH_SIZE)\n    print(f'  Train: {train_emb.shape}')\n\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\n    print('\\nGenerating standalone submission...')\n    img_map = {n: i for i, n in enumerate(test_names)}\n    generate_submission(test_df, test_emb, img_map, OUT_DIR / f'submission_{VERSION}.csv')\n    generate_submission(test_df, test_emb, img_map, Path('/kaggle/working/submission.csv'))\n\n    print(f'\\n{\"=\"*60}')\n    print(f'{VERSION} complete! Total: {(time.time()-t0)/60:.1f} min')\n    print(f'{\"=\"*60}')","metadata":{"_uuid":"38f8f463-4429-4f2c-8d98-02cfd7dd50eb","_cell_guid":"ab7eeda5-5640-4d53-b418-fc27b0ab7c64","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null}]}