{"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":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install timm -q","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:49.275097Z","iopub.execute_input":"2026-02-25T07:38:49.275452Z","iopub.status.idle":"2026-02-25T07:38:52.653093Z","shell.execute_reply.started":"2026-02-25T07:38:49.275425Z","shell.execute_reply":"2026-02-25T07:38:52.652113Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\nimport timm\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.decomposition import PCA\n\nimport torchvision.transforms as transforms\nfrom PIL import Image\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"Device:\", device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:52.654915Z","iopub.execute_input":"2026-02-25T07:38:52.655252Z","iopub.status.idle":"2026-02-25T07:38:52.661667Z","shell.execute_reply.started":"2026-02-25T07:38:52.655225Z","shell.execute_reply":"2026-02-25T07:38:52.660890Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"DATA_PATH = \"/kaggle/input/jaguar-re-id\"\n\ntrain_df = pd.read_csv(f\"{DATA_PATH}/train.csv\")\ntest_df = pd.read_csv(f\"{DATA_PATH}/test.csv\")\n\nprint(\"Train shape:\", train_df.shape)\nprint(\"Test pairs:\", test_df.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:52.662610Z","iopub.execute_input":"2026-02-25T07:38:52.662851Z","iopub.status.idle":"2026-02-25T07:38:52.742817Z","shell.execute_reply.started":"2026-02-25T07:38:52.662831Z","shell.execute_reply":"2026-02-25T07:38:52.742233Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TRAIN_IMAGE_PATH = os.path.join(DATA_PATH, \"train\")\nif os.path.isdir(os.path.join(TRAIN_IMAGE_PATH, \"train\")):\n    TRAIN_IMAGE_PATH = os.path.join(TRAIN_IMAGE_PATH, \"train\")\n\nTEST_IMAGE_PATH = os.path.join(DATA_PATH, \"test\")\nif os.path.isdir(os.path.join(TEST_IMAGE_PATH, \"test\")):\n    TEST_IMAGE_PATH = os.path.join(TEST_IMAGE_PATH, \"test\")\n\nprint(TRAIN_IMAGE_PATH)\nprint(TEST_IMAGE_PATH)\n\nprint(\"Sample image exists:\",\n      os.path.exists(os.path.join(TRAIN_IMAGE_PATH, train_df.iloc[0][\"filename\"])))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:52.744231Z","iopub.execute_input":"2026-02-25T07:38:52.744466Z","iopub.status.idle":"2026-02-25T07:38:52.755990Z","shell.execute_reply.started":"2026-02-25T07:38:52.744445Z","shell.execute_reply":"2026-02-25T07:38:52.755424Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"le = LabelEncoder()\ntrain_df[\"label\"] = le.fit_transform(train_df[\"ground_truth\"])\nnum_classes = train_df[\"label\"].nunique()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:52.756951Z","iopub.execute_input":"2026-02-25T07:38:52.757272Z","iopub.status.idle":"2026-02-25T07:38:52.763059Z","shell.execute_reply.started":"2026-02-25T07:38:52.757251Z","shell.execute_reply":"2026-02-25T07:38:52.762180Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sample_imgs = train_df[\"filename\"].sample(6).values\n\nplt.figure(figsize=(12,6))\nfor i, img_name in enumerate(sample_imgs):\n    img = Image.open(os.path.join(TRAIN_IMAGE_PATH, img_name))\n    plt.subplot(2,3,i+1)\n    plt.imshow(img)\n    plt.title(img_name)\n    plt.axis(\"off\")\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:52.764082Z","iopub.execute_input":"2026-02-25T07:38:52.764461Z","iopub.status.idle":"2026-02-25T07:38:59.623589Z","shell.execute_reply.started":"2026-02-25T07:38:52.764424Z","shell.execute_reply":"2026-02-25T07:38:59.622706Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_tfms = transforms.Compose([\n    transforms.Resize((384,384)),\n    transforms.RandomHorizontalFlip(),\n    transforms.ColorJitter(0.2,0.2,0.2,0.1),\n\n    transforms.ToTensor(),  # MUST come before RandomErasing\n\n    transforms.RandomErasing(p=0.3),\n\n    transforms.Normalize([0.485,0.456,0.406],\n                         [0.229,0.224,0.225])\n])\n\nvalid_tfms = transforms.Compose([\n    transforms.Resize((384,384)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485,0.456,0.406],\n                         [0.229,0.224,0.225])\n])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:59.624667Z","iopub.execute_input":"2026-02-25T07:38:59.624903Z","iopub.status.idle":"2026-02-25T07:38:59.630161Z","shell.execute_reply.started":"2026-02-25T07:38:59.624882Z","shell.execute_reply":"2026-02-25T07:38:59.629435Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class JaguarDataset(Dataset):\n    def __init__(self, df, img_dir, transforms=None):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = img_dir\n        self.transforms = transforms\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img = Image.open(os.path.join(self.img_dir, row[\"filename\"])).convert(\"RGB\")\n        if self.transforms:\n            img = self.transforms(img)\n        return img, row[\"label\"]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:59.630939Z","iopub.execute_input":"2026-02-25T07:38:59.631175Z","iopub.status.idle":"2026-02-25T07:38:59.643911Z","shell.execute_reply.started":"2026-02-25T07:38:59.631156Z","shell.execute_reply":"2026-02-25T07:38:59.643288Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ArcFace(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, x, labels):\n        cosine = F.linear(F.normalize(x), F.normalize(self.weight))\n        theta = torch.acos(torch.clamp(cosine, -1+1e-7, 1-1e-7))\n        target_logits = torch.cos(theta + self.m)\n\n        one_hot = F.one_hot(labels, cosine.size(1)).float()\n        output = cosine * (1 - one_hot) + target_logits * one_hot\n        output *= self.s\n        return output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:59.644706Z","iopub.execute_input":"2026-02-25T07:38:59.644950Z","iopub.status.idle":"2026-02-25T07:38:59.657420Z","shell.execute_reply.started":"2026-02-25T07:38:59.644930Z","shell.execute_reply":"2026-02-25T07:38:59.656800Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BatchHardTripletLoss(nn.Module):\n    def __init__(self, margin=0.3):\n        super().__init__()\n        self.margin = margin\n\n    def forward(self, embeddings, labels):\n        dist_matrix = torch.cdist(embeddings, embeddings)\n        labels = labels.unsqueeze(1)\n\n        mask_pos = labels.eq(labels.T)\n        mask_neg = ~mask_pos\n\n        hardest_pos = (dist_matrix * mask_pos.float()).max(1)[0]\n        hardest_neg = (dist_matrix + 1e5 * mask_pos.float()).min(1)[0]\n\n        loss = F.relu(hardest_pos - hardest_neg + self.margin)\n        return loss.mean()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:59.659181Z","iopub.execute_input":"2026-02-25T07:38:59.659459Z","iopub.status.idle":"2026-02-25T07:38:59.669313Z","shell.execute_reply.started":"2026-02-25T07:38:59.659439Z","shell.execute_reply":"2026-02-25T07:38:59.668436Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class JaguarModel(nn.Module):\n    def __init__(self, num_classes):\n        super().__init__()\n        self.backbone = timm.create_model(\"convnext_base\", pretrained=True, num_classes=0)\n        self.embedding = nn.Linear(self.backbone.num_features, 512)\n        self.bn = nn.BatchNorm1d(512)\n        self.arcface = ArcFace(512, num_classes)\n\n    def forward(self, x, labels=None):\n        features = self.backbone(x)\n        emb = self.embedding(features)\n        emb = self.bn(emb)\n        emb = F.normalize(emb)\n\n        if labels is not None:\n            logits = self.arcface(emb, labels)\n            return logits, emb\n        else:\n            return emb","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:59.670682Z","iopub.execute_input":"2026-02-25T07:38:59.670947Z","iopub.status.idle":"2026-02-25T07:38:59.685954Z","shell.execute_reply.started":"2026-02-25T07:38:59.670917Z","shell.execute_reply":"2026-02-25T07:38:59.685292Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_dataset = JaguarDataset(train_df, TRAIN_IMAGE_PATH, train_tfms)\ntrain_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:59.686888Z","iopub.execute_input":"2026-02-25T07:38:59.687191Z","iopub.status.idle":"2026-02-25T07:38:59.700849Z","shell.execute_reply.started":"2026-02-25T07:38:59.687161Z","shell.execute_reply":"2026-02-25T07:38:59.700118Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = JaguarModel(num_classes).to(device)\n\nce_loss = nn.CrossEntropyLoss()\ntriplet_loss = BatchHardTripletLoss(margin=0.3)\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:38:59.702680Z","iopub.execute_input":"2026-02-25T07:38:59.702898Z","iopub.status.idle":"2026-02-25T07:39:01.151476Z","shell.execute_reply.started":"2026-02-25T07:38:59.702879Z","shell.execute_reply":"2026-02-25T07:39:01.150821Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"scaler = torch.amp.GradScaler(\"cuda\")\n\ntotal_losses, ce_losses, tri_losses, lr_history = [], [], [], []\n\nfor epoch in range(20):\n    model.train()\n    epoch_total = 0\n    epoch_ce = 0\n    epoch_tri = 0\n\n    for images, labels in tqdm(train_loader):\n        images = images.to(device)\n        labels = labels.to(device)\n\n        optimizer.zero_grad()\n\n        with torch.amp.autocast(device_type=\"cuda\"):\n            logits, embeddings = model(images, labels)\n            loss_ce = ce_loss(logits, labels)\n            loss_tri = triplet_loss(embeddings, labels)\n            loss = loss_ce + loss_tri\n\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n\n        epoch_total += loss.item()\n        epoch_ce += loss_ce.item()\n        epoch_tri += loss_tri.item()\n\n    scheduler.step()\n\n    total_losses.append(epoch_total/len(train_loader))\n    ce_losses.append(epoch_ce/len(train_loader))\n    tri_losses.append(epoch_tri/len(train_loader))\n    lr_history.append(optimizer.param_groups[0]['lr'])\n\n    print(f\"Epoch {epoch} | Total: {total_losses[-1]:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T07:39:01.152491Z","iopub.execute_input":"2026-02-25T07:39:01.152707Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.figure(figsize=(12,5))\n\nplt.subplot(1,2,1)\nplt.plot(total_losses, label=\"Total\")\nplt.plot(ce_losses, label=\"CE\")\nplt.plot(tri_losses, label=\"Triplet\")\nplt.legend()\nplt.title(\"Loss Components\")\nplt.grid()\n\nplt.subplot(1,2,2)\nplt.plot(lr_history)\nplt.title(\"Learning Rate\")\nplt.grid()\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.eval()\n\ntest_images = sorted(os.listdir(TEST_IMAGE_PATH))\nembeddings = []\n\nwith torch.no_grad():\n    for img_name in tqdm(test_images):\n        img = Image.open(os.path.join(TEST_IMAGE_PATH, img_name)).convert(\"RGB\")\n        img = valid_tfms(img).unsqueeze(0).to(device)\n        emb = model(img)\n        embeddings.append(emb.cpu())\n\nembeddings = torch.cat(embeddings)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"similarity_matrix = torch.mm(embeddings, embeddings.T)\n\nplt.figure(figsize=(8,6))\nplt.imshow(similarity_matrix[:50,:50].numpy(), cmap=\"viridis\")\nplt.colorbar()\nplt.title(\"Similarity Heatmap\")\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pca = PCA(n_components=2)\nemb_2d = pca.fit_transform(embeddings[:300])\n\nplt.figure(figsize=(6,6))\nplt.scatter(emb_2d[:,0], emb_2d[:,1], s=10)\nplt.title(\"Embedding Distribution (PCA)\")\nplt.grid()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission = []\n\nfor _, row in test_df.iterrows():\n    q = test_images.index(row[\"query_image\"])\n    g = test_images.index(row[\"gallery_image\"])\n    submission.append(similarity_matrix[q,g].item())\n\nsub_df = pd.DataFrame({\n    \"row_id\": test_df[\"row_id\"],\n    \"similarity\": submission\n})\n\nsub_df.to_csv(\"submission.csv\", index=False)\nprint(\"Submission saved!\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}