{"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":"\"\"\"\nv07c: ConvNeXt-Large @ 384px — Model C (cross-architecture diversity)\n=====================================================================\nPhase 1 of the v07 ensemble pipeline.\n\nKey differences from v07a/v07b (EVA-02):\n  - CNN backbone (local inductive bias) vs ViT (global attention)\n  - 384px (vs 448px) — ConvNeXt uses 7×7 stem, no patch size constraint\n  - Native (B, C, H, W) features — no ViT token reshape needed\n  - Differential LR: backbone 5e-5, head 1e-3 (proven in v04)\n  - 15 epochs (CNNs need more iterations to converge)\n  - Embedding dim: 1536 (ConvNeXt-Large native)\n\nOutputs: model_v07c.pth, embeddings_v07c_{train,test}.npy, filenames, submission\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 = 'v07c'\nSEED = 42\n\nMODEL_NAME = 'convnext_large.fb_in22k_ft_in1k_384'\nIMG_SIZE = 384\nNUM_CLASSES = 31\n\nNUM_EPOCHS = 15\nBATCH_SIZE = 8\nGRAD_ACCUM = 2        # effective batch = 16\nLR_BACKBONE = 5e-5\nLR_HEAD = 1e-3\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 ───────────────────────────────────────────────────────────\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    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 ConvNeXtReIDModel(nn.Module):\n    \"\"\"ConvNeXt-Large backbone with GeM pooling and ArcFace head.\n\n    ConvNeXt natively outputs (B, C, H, W) feature maps — no reshape needed.\n    With 384px input and 7×7 stem + 4 downsamples, output is 12×12 spatial.\"\"\"\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        # ConvNeXt: already (B, C, H, W) — but some timm versions return (B, H*W, C)\n        if features.dim() == 3:\n            B, N, C = features.shape\n            H = W = int(math.sqrt(N))\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=8):\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 = ConvNeXtReIDModel().to(DEVICE)\n\n    # Gradient checkpointing (if available for this backbone)\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    print(f'Model: {MODEL_NAME}')\n    print(f'Parameters: {n_params:,} | Embedding dim: {model.feat_dim}')\n\n    # Differential learning rate: backbone learns slower, head learns faster\n    backbone_params = list(model.backbone.parameters())\n    head_params = (list(model.gem.parameters()) +\n                   list(model.bn.parameters()) +\n                   list(model.head.parameters()))\n    optimizer = torch.optim.AdamW([\n        {'params': backbone_params, 'lr': LR_BACKBONE},\n        {'params': head_params, 'lr': LR_HEAD},\n    ], weight_decay=WEIGHT_DECAY)\n\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')\n    print(f'Epochs: {NUM_EPOCHS} | Batch: {BATCH_SIZE}×{GRAD_ACCUM}')\n    print(f'LR: backbone={LR_BACKBONE}, head={LR_HEAD}')\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        lr_bb = optimizer.param_groups[0]['lr']\n        lr_hd = optimizer.param_groups[1]['lr']\n        print(f'Epoch {epoch+1:2d}/{NUM_EPOCHS} | '\n              f'Loss: {loss:.4f} | LR: bb={lr_bb:.2e} hd={lr_hd:.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":"c4ebc162-8e22-4c4e-8ac3-a9a4400ee7eb","_cell_guid":"58ae9c5f-9d2f-4025-b4d4-c2264bf94ec7","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null}]}