{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":126777,"databundleVersionId":15314950}],"dockerImageVersionId":31286,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\n\nimport os\nimport matplotlib.pyplot as plt\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\n\nimport timm\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:14.123899Z","iopub.execute_input":"2026-03-05T03:10:14.124256Z","iopub.status.idle":"2026-03-05T03:10:27.421833Z","shell.execute_reply.started":"2026-03-05T03:10:14.124222Z","shell.execute_reply":"2026-03-05T03:10:27.421077Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nclass Config:\n    seed = 42\n    model_name = \"hf-hub:BVRA/MegaDescriptor-L-384\"\n    img_size = 384\n    embedding_dim = 1536\n    num_classes = 31\n    num_epochs = 25\n    batch_size = 2\n    grad_accum = 8\n    lr = 2e-4\n    weight_decay = 1e-3\n    arcface_s = 64.0\n    arcface_m = 0.35\n    use_tta = True\n    use_qe = True\n    use_rerank = True\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    device_type = \"cuda\" if torch.cuda.is_available() else \"cpu\"\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(Config.seed)\nprint(\"Device:\", Config.device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:27.423373Z","iopub.execute_input":"2026-03-05T03:10:27.423780Z","iopub.status.idle":"2026-03-05T03:10:27.691157Z","shell.execute_reply.started":"2026-03-05T03:10:27.423756Z","shell.execute_reply":"2026-03-05T03:10:27.690334Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(torch.cuda.is_available())        \nprint(torch.cuda.get_device_name(0))    \nprint(Config.device) ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:27.692194Z","iopub.execute_input":"2026-03-05T03:10:27.692455Z","iopub.status.idle":"2026-03-05T03:10:27.744410Z","shell.execute_reply.started":"2026-03-05T03:10:27.692432Z","shell.execute_reply":"2026-03-05T03:10:27.743631Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BASE      = Path('/kaggle/input/competitions/jaguar-re-id')\nTRAIN_DIR = BASE / 'train' / 'train'\nTEST_DIR  = BASE / 'test' / 'test'\n\ntrain_df = pd.read_csv(BASE / 'train.csv')\ntest_df  = pd.read_csv(BASE / 'test.csv')\n\nprint('train_df shape:', train_df.shape)\nprint('test_df  shape:', test_df.shape)\nprint(train_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:27.745415Z","iopub.execute_input":"2026-03-05T03:10:27.745684Z","iopub.status.idle":"2026-03-05T03:10:27.891267Z","shell.execute_reply.started":"2026-03-05T03:10:27.745663Z","shell.execute_reply":"2026-03-05T03:10:27.890613Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig,axes=plt.subplots(3,2,figsize=(10,12))\nfig.suptitle('Original vs Alpha Mask Applied', fontsize=14, fontweight='bold')\nfor i in range(3):\n    img_path=TRAIN_DIR/train_df['filename'].iloc[i]\n    raw=Image.open(img_path).convert('RGB')\n    axes[i][0].imshow(raw)\n    axes[i][0].set_title(f'Original — {train_df[\"ground_truth\"].iloc[i]}')\n    axes[i][0].axis(\"off\")\n    rgba=Image.open(img_path)\n    r,g,b,alpha=rgba.split()\n    background = Image.new(\"RGB\", rgba.size, (0, 0, 0))\n    masked = Image.composite(rgba.convert(\"RGB\"), background, alpha)\n    \n    axes[i][1].imshow(masked)\n    axes[i][1].set_title(f'Masked — {train_df[\"ground_truth\"].iloc[i]}')\n    axes[i][1].axis('off')\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:27.892495Z","iopub.execute_input":"2026-03-05T03:10:27.893148Z","iopub.status.idle":"2026-03-05T03:10:30.665004Z","shell.execute_reply.started":"2026-03-05T03:10:27.893114Z","shell.execute_reply":"2026-03-05T03:10:30.664287Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class JaguarDataset(Dataset):\n    def __init__(self, df, img_dir, transform=None, is_test=False):\n        self.df = df.copy()\n        self.img_dir = Path(img_dir)\n        self.transform = transform\n        self.is_test = is_test\n\n        if not is_test:\n            unique_ids = sorted(self.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 load_image(self, img_path):\n        img = Image.open(img_path)\n        if img.mode == \"RGBA\":\n            r, g, b, alpha = img.split()\n            background = Image.new(\"RGB\", img.size, (0, 0, 0))\n            return Image.composite(img.convert(\"RGB\"), background, alpha)\n        else:\n            return img.convert(\"RGB\")\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = self.img_dir / row[\"filename\"]\n        try:\n            img = self.load_image(img_path)\n        except Exception as e:\n            print(f\"Warning: could not load {img_path}: {e}\")\n            img = Image.new(\"RGB\", (Config.img_size, Config.img_size), (0, 0, 0))\n\n        if self.transform:\n            img = self.transform(img)\n\n        if self.is_test:\n            return img, row[\"filename\"]\n        return img, torch.tensor(row[\"label\"], dtype=torch.long)\n\n    def get_sample_weights(self):\n        class_counts = self.df[\"ground_truth\"].value_counts()\n        weights = self.df[\"ground_truth\"].map(lambda x: 1.0 / class_counts[x])\n        return weights.values\n\nprint(\"The jaguar dataset\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:30.666210Z","iopub.execute_input":"2026-03-05T03:10:30.666550Z","iopub.status.idle":"2026-03-05T03:10:30.676120Z","shell.execute_reply.started":"2026-03-05T03:10:30.666523Z","shell.execute_reply":"2026-03-05T03:10:30.675379Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_transform = transforms.Compose([\n    transforms.Resize((Config.img_size, Config.img_size)),\n    transforms.RandomHorizontalFlip(p=0.5),\n    transforms.RandomVerticalFlip(p=0.1),\n    transforms.RandomAffine(\n        degrees=30,\n        translate=(0.1, 0.1),\n        scale=(0.85, 1.15),\n        shear=10\n    ),\n    transforms.RandomPerspective(distortion_scale=0.2, p=0.3),\n    transforms.ColorJitter(\n        brightness=0.3,\n        contrast=0.3,\n        saturation=0.2,\n        hue=0.05\n    ),\n    transforms.RandomGrayscale(p=0.05),\n    transforms.ToTensor(),\n    transforms.Normalize([0.481, 0.457, 0.408], [0.268, 0.261, 0.275]),\n    transforms.RandomErasing(p=0.3, scale=(0.02, 0.2))\n])\n\ntest_transform = transforms.Compose([\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])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:30.679456Z","iopub.execute_input":"2026-03-05T03:10:30.679798Z","iopub.status.idle":"2026-03-05T03:10:30.833031Z","shell.execute_reply.started":"2026-03-05T03:10:30.679769Z","shell.execute_reply":"2026-03-05T03:10:30.832175Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.utils.data import WeightedRandomSampler\ntrain_dataset=JaguarDataset(train_df,TRAIN_DIR,train_transform)\nsample_weights=train_dataset.get_sample_weights()\nsampler=WeightedRandomSampler(\n    weights=sample_weights,\n    num_samples=len(sample_weights),\n    replacement=True\n)\ntrain_loader=DataLoader(\n    train_dataset,\n    batch_size=Config.batch_size,\n    sampler=sampler,\n    num_workers=2,\n    pin_memory=True,\n    drop_last=True \n)\n\n\nprint(f\"The total number of training samples are :{len(train_dataset)}\")\nprint(f\"Bathces per epochs are:{len(train_loader)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:30.834206Z","iopub.execute_input":"2026-03-05T03:10:30.834632Z","iopub.status.idle":"2026-03-05T03:10:30.880506Z","shell.execute_reply.started":"2026-03-05T03:10:30.834594Z","shell.execute_reply":"2026-03-05T03:10:30.879755Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install -q wildlife-tools","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:30.881440Z","iopub.execute_input":"2026-03-05T03:10:30.881769Z","iopub.status.idle":"2026-03-05T03:10:36.344070Z","shell.execute_reply.started":"2026-03-05T03:10:30.881746Z","shell.execute_reply":"2026-03-05T03:10:36.343210Z"}},"outputs":[],"execution_count":null},{"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        if x.dim() == 2:\n            return 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=64.0, m=0.35):\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 JaguarModel(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,\n            s=Config.arcface_s, m=Config.arcface_m\n        )\n\n    def forward(self, x, label=None):\n        features = model.backbone.forward_features(x)\n    \n        if features.dim() == 3:\n            features = features.mean(dim=1)\n        elif features.dim() == 4:\n           \n            features = features.mean(dim=[1, 2])  \n        \n        emb = self.bn(features)\n        if label is not None:\n            return self.head(emb, label)\n        return emb\n\n\nmodel = JaguarModel().to(Config.device)\nprint(\"Model has been done \")\nprint(f\"Feature dim: {model.feat_dim}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:36.345254Z","iopub.execute_input":"2026-03-05T03:10:36.345581Z","iopub.status.idle":"2026-03-05T03:10:46.332276Z","shell.execute_reply.started":"2026-03-05T03:10:36.345544Z","shell.execute_reply":"2026-03-05T03:10:46.331579Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(model.bn)  \nprint(type(model))\nmodel = model.to(Config.device)\nprint(next(model.parameters()).device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:46.333193Z","iopub.execute_input":"2026-03-05T03:10:46.333454Z","iopub.status.idle":"2026-03-05T03:10:46.343297Z","shell.execute_reply.started":"2026-03-05T03:10:46.333432Z","shell.execute_reply":"2026-03-05T03:10:46.342549Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc\ntorch.cuda.empty_cache()\ngc.collect()\nprint(f\"GPU memory free: {torch.cuda.mem_get_info()[0]/1e9:.2f} GB\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:46.344140Z","iopub.execute_input":"2026-03-05T03:10:46.344409Z","iopub.status.idle":"2026-03-05T03:10:46.700996Z","shell.execute_reply.started":"2026-03-05T03:10:46.344388Z","shell.execute_reply":"2026-03-05T03:10:46.700245Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_epoch(model, loader, optimizer, criterion, scaler, scheduler):\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        with torch.amp.autocast(Config.device_type):\n            loss = criterion(model(imgs, labels), labels)\n            loss = loss / Config.grad_accum\n        scaler.scale(loss).backward()\n        if (i + 1) % Config.grad_accum == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n            scheduler.step()  # ✅ moved inside, per batch\n        loss_meter += loss.item() * Config.grad_accum\n    return loss_meter / len(loader)\n\n\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.OneCycleLR(\n    optimizer,\n    max_lr=Config.lr,\n    epochs=Config.num_epochs,\n    steps_per_epoch=len(train_loader)\n)\n\nprint(\"Training starting...\")\nfor epoch in range(Config.num_epochs):\n    loss = train_epoch(model, train_loader, optimizer, nn.CrossEntropyLoss(), scaler, scheduler)  \n    print(f\"Epoch {epoch+1}/{Config.num_epochs} | Loss: {loss:.4f} | LR: {optimizer.param_groups[0]['lr']:.2e}\")\n\nprint(\"Training done\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T03:10:46.701972Z","iopub.execute_input":"2026-03-05T03:10:46.702270Z","iopub.status.idle":"2026-03-05T05:58:21.622217Z","shell.execute_reply.started":"2026-03-05T03:10:46.702248Z","shell.execute_reply":"2026-03-05T05:58:21.621393Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import inspect\nprint(inspect.getsource(model.forward))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T05:58:21.623929Z","iopub.execute_input":"2026-03-05T05:58:21.624256Z","iopub.status.idle":"2026-03-05T05:58:21.630289Z","shell.execute_reply.started":"2026-03-05T05:58:21.624219Z","shell.execute_reply":"2026-03-05T05:58:21.629693Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.eval()\nimgs, labels = next(iter(train_loader))\nimgs = imgs.to(Config.device)\n\nwith torch.no_grad():\n    features = model.backbone.forward_features(imgs)\n    print(\"features shape:\", features.shape)\n    print(\"features dim:\", features.dim())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T05:58:21.631223Z","iopub.execute_input":"2026-03-05T05:58:21.631465Z","iopub.status.idle":"2026-03-05T05:58:23.477021Z","shell.execute_reply.started":"2026-03-05T05:58:21.631446Z","shell.execute_reply":"2026-03-05T05:58:23.476228Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\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\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}),\n        TEST_DIR, test_transform, is_test=True\n    ),\n    batch_size=Config.batch_size,\n    shuffle=False,\n    num_workers=0,\n    pin_memory=True\n)\n\nemb, names = extract_features(model, test_loader)\nprint(f\"Extracted {len(names)} embeddings of shape {emb.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T05:58:23.478294Z","iopub.execute_input":"2026-03-05T05:58:23.478523Z","iopub.status.idle":"2026-03-05T06:01:42.497730Z","shell.execute_reply.started":"2026-03-05T05:58:23.478500Z","shell.execute_reply":"2026-03-05T06:01:42.497117Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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=10, k2=4, lambda_value=0.3):\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))\n\n\nif Config.use_qe:\n    emb = query_expansion(emb)\n\nsim_matrix = emb @ emb.T\n\nif Config.use_rerank:\n    sim_matrix = k_reciprocal_rerank(sim_matrix)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T06:05:29.395519Z","iopub.execute_input":"2026-03-05T06:05:29.395806Z","iopub.status.idle":"2026-03-05T06:05:29.699773Z","shell.execute_reply.started":"2026-03-05T06:05:29.395780Z","shell.execute_reply":"2026-03-05T06:05:29.699193Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"img_map = {n: i for i, n in enumerate(names)}\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(float(np.clip(s, 0.0, 1.0)))\n\nsub = pd.DataFrame({\"row_id\": test_df[\"row_id\"], \"similarity\": preds})\nsub.to_csv(\"submission.csv\", index=False)\n\nprint(sub.head())\nprint(len(sub))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-05T06:06:00.944250Z","iopub.execute_input":"2026-03-05T06:06:00.944542Z","iopub.status.idle":"2026-03-05T06:06:07.490783Z","shell.execute_reply.started":"2026-03-05T06:06:00.944517Z","shell.execute_reply":"2026-03-05T06:06:07.490163Z"}},"outputs":[],"execution_count":null}]}