{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":126777,"databundleVersionId":15314950}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -qU timm","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Imports\n","metadata":{}},{"cell_type":"code","source":"import os\nimport math\nimport random\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nfrom PIL import Image\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","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Config\n","metadata":{}},{"cell_type":"code","source":"class Config:\n    seed = 42\n    model_name = \"eva02_large_patch14_448.mim_m38m_ft_in22k_in1k\"\n\n    img_size = 448\n    embedding_dim = 1024\n    num_classes = 31\n\n    num_epochs = 10\n    batch_size = 4\n    grad_accum = 4\n\n    lr = 2e-5\n    weight_decay = 1e-3\n\n    arcface_s = 30.0\n    arcface_m = 0.50\n\n    use_tta = True\n    use_qe = True\n    use_rerank = True\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    device_type = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\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\n\nseed_everything(Config.seed)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Dataset\n","metadata":{}},{"cell_type":"code","source":"class JaguarDataset(Dataset):\n    def __init__(self, df, img_dir, transform=None, is_test=False):\n        self.df = df\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:\n            img = Image.new(\"RGB\", (Config.img_size, Config.img_size))\n\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)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Transforms\n","metadata":{}},{"cell_type":"code","source":"train_transform = transforms.Compose(\n    [\n        transforms.Resize((Config.img_size, Config.img_size)),\n        transforms.RandomHorizontalFlip(),\n        transforms.RandomAffine(degrees=15, translate=(0.1, 0.1), scale=(0.9, 1.1)),\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([0.481, 0.457, 0.408], [0.268, 0.261, 0.275]),\n        transforms.RandomErasing(p=0.25),\n    ]\n)\n\ntest_transform = transforms.Compose(\n    [\n        transforms.Resize((Config.img_size, Config.img_size)),\n        transforms.ToTensor(),\n        transforms.Normalize([0.481, 0.457, 0.408], [0.268, 0.261, 0.275]),\n    ]\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model\n","metadata":{}},{"cell_type":"code","source":"class GeM(nn.Module):\n    def __init__(self, p=3, eps=1e-6):\n        super(GeM, self).__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), (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 EVABoss(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.backbone = timm.create_model(\n            Config.model_name, pretrained=True, num_classes=0\n        )\n        self.feat_dim = self.backbone.num_features\n        self.gem = GeM()\n        self.bn = nn.BatchNorm1d(self.feat_dim)\n        self.head = ArcFaceLayer(\n            self.feat_dim, Config.num_classes, s=Config.arcface_s, m=Config.arcface_m\n        )\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\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","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Utils\n","metadata":{}},{"cell_type":"code","source":"def train_epoch(model, loader, optimizer, criterion, scaler):\n    model.train()\n    loss_meter = 0\n    for i, (imgs, labels) in enumerate(tqdm(loader, leave=False)):\n        imgs, labels = imgs.to(Config.device), labels.to(Config.device)\n\n        with torch.amp.autocast(Config.device_type):\n            loss = criterion(model(imgs, labels), labels)\n            loss = loss / Config.grad_accum\n\n        scaler.scale(loss).backward()\n\n        if (i + 1) % Config.grad_accum == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n\n        loss_meter += loss.item() * Config.grad_accum\n    return loss_meter / len(loader)\n\n\n@torch.no_grad()\ndef extract_features(model, loader):\n    model.eval()\n    feats, names = [], []\n    for imgs, fnames in tqdm(loader, desc=\"Inference\"):\n        imgs = imgs.to(Config.device)\n        f1 = model(imgs)\n        if Config.use_tta:\n            f2 = model(torch.flip(imgs, [3]))\n            f1 = (f1 + f2) / 2\n        feats.append(F.normalize(f1, dim=1).cpu())\n        names.extend(fnames)\n    return torch.cat(feats, dim=0).numpy(), names\n\n\ndef query_expansion(emb, top_k=3):\n    print(\"Applying QE...\")\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(\"Applying Re-ranking...\")\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 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))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Execution\n","metadata":{}},{"cell_type":"code","source":"TRAIN_CSV = \"/kaggle/input/jaguar-re-id/train.csv\"\nTEST_CSV = \"/kaggle/input/jaguar-re-id/test.csv\"\nTRAIN_DIR = \"/kaggle/input/jaguar-re-id/train/train\"\nTEST_DIR = \"/kaggle/input/jaguar-re-id/test/test\"\n\ntrain_df = pd.read_csv(TRAIN_CSV)\ntest_df = pd.read_csv(TEST_CSV)\n\ntrain_loader = DataLoader(\n    JaguarDataset(train_df, TRAIN_DIR, train_transform),\n    batch_size=Config.batch_size,\n    shuffle=True,\n    num_workers=2,\n    pin_memory=False,\n)\nmodel = EVABoss().to(Config.device)\noptimizer = torch.optim.AdamW(\n    model.parameters(), lr=Config.lr, weight_decay=Config.weight_decay\n)\nscaler = torch.amp.GradScaler(Config.device_type)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n    optimizer, T_max=Config.num_epochs\n)\n\nprint(f\"🔥 Training EVA-02 Large (448px)...\")\n\nfor epoch in range(Config.num_epochs):\n    loss = train_epoch(model, train_loader, optimizer, nn.CrossEntropyLoss(), scaler)\n    scheduler.step()\n    print(\n        f\"Epoch {epoch+1}/{Config.num_epochs} | Loss: {loss:.4f} | LR: {scheduler.get_last_lr()[0]:.2e}\"\n    )\n\ntorch.save(model.state_dict(), 'model_weights.pth')\n\nunique_test = sorted(set(test_df[\"query_image\"]) | set(test_df[\"gallery_image\"]))\ntest_loader = DataLoader(\n    JaguarDataset(\n        pd.DataFrame({\"filename\": unique_test}), TEST_DIR, test_transform, True\n    ),\n    batch_size=Config.batch_size,\n    shuffle=False,\n    num_workers=2,\n    pin_memory=False,\n)\n\nemb, names = extract_features(model, test_loader)\nimg_map = {n: i for i, n in enumerate(names)}\n\nif Config.use_qe:\n    emb = query_expansion(emb)\nsim_matrix = emb @ emb.T\nif Config.use_rerank:\n    sim_matrix = k_reciprocal_rerank(sim_matrix)\n\npreds = []\nfor _, 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, s)))\n\nsub = pd.DataFrame({\"row_id\": test_df[\"row_id\"], \"similarity\": preds})\nsub.to_csv(\"submission.csv\", index=False)\nprint(f\"✅ Done! Mean Sim: {np.mean(preds):.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}