{"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":[{"sourceId":126777,"databundleVersionId":15314950,"sourceType":"competition"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Jaguar Re-ID: Vision Transformer & ArcFace Pipeline\n\n### Overview\nThis notebook provides a concise and powerful pipeline for Jaguar Re-Identification using **DINOv2 (Vision Transformer)** and **ArcFace**. It is optimized to achieve high mAP by focusing on fine-grained rosette pattern recognition.\n\n### Core Strategy\n* **Backbone**: **vit_base_patch14_dinov2**(Pretrained ViT)\n* **Metric Learning**: **ArcFace** for superior identity separation.\n* **Input Size**: 518x518 to capture high-resolution spot details.\n* **Learning**: Cosine Annealing LR & Label Smoothing for optimal convergence.\n* **Inference**: Horizontal Flip TTA to stabilize and boost LB score.\n\n### Performance Checklist\n- High-resolution processing (518px)\n- Test-Time Augmentation (TTA)\n- Normalized Feature Embeddings","metadata":{}},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\nimport torchvision.transforms as T\nimport timm\nfrom sklearn.preprocessing import LabelEncoder\nfrom tqdm import tqdm\n\n# ====================================================\n# Configuration\n# ====================================================\nINPUT_DIR = \"/kaggle/input/jaguar-re-id\"\n\nCFG = {\n    'train_csv': os.path.join(INPUT_DIR, \"train.csv\"),\n    'test_csv': os.path.join(INPUT_DIR, \"test.csv\"),\n    'sample_sub': os.path.join(INPUT_DIR, \"sample_submission.csv\"),\n    'train_dir': os.path.join(INPUT_DIR, \"train/train\"),\n    'test_dir': os.path.join(INPUT_DIR, \"test/test\"),\n    \n    'model_name': 'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k', \n    'img_size': 448,\n    'batch_size': 4,\n    'epochs': 15,\n    'embedding_size': 1024,\n    'lr': 8e-6,\n    'weight_decay': 1e-4,\n    's': 32.0,\n    'm': 0.50,\n    'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n}\n\n# ====================================================\n# Define Dataset\n# ====================================================\nclass JaguarDataset(Dataset):\n    def __init__(self, df, img_dir, transform=None, is_test=False):\n        self.df = df\n        self.img_dir = img_dir\n        self.transform = transform\n        self.is_test = is_test\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        filename = self.df.iloc[idx]['filename']\n        img_path = os.path.join(self.img_dir, filename)\n        \n        try:\n            image = Image.open(img_path).convert('RGB')\n        except FileNotFoundError:\n            raise FileNotFoundError(f\"cannot find images: {img_path}\")\n            \n        if self.transform:\n            image = self.transform(image)\n        \n        if self.is_test:\n            return image, filename\n        \n        label = torch.tensor(self.df.iloc[idx]['label'], dtype=torch.long)\n        return image, label","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ====================================================\n# Define the model(ViT + ArcFace)\n# ====================================================\nclass ArcMarginProduct(nn.Module):\n    def __init__(self, in_features, out_features, s=64.0, m=0.50):\n        super().__init__()\n        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))\n        nn.init.xavier_uniform_(self.weight)\n        self.s = s\n        self.m = m\n\n    def forward(self, input, label):\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))\n        phi = cosine * np.cos(self.m) - sine * np.sin(self.m)\n        one_hot = torch.zeros(cosine.size(), device=CFG['device'])\n        one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n        return output * self.s\n\nclass JaguarModel(nn.Module):\n    def __init__(self, n_classes):\n        super().__init__()\n        self.backbone = timm.create_model(\n            CFG['model_name'], \n            pretrained=True, \n            num_classes=0,\n            dynamic_img_size=True \n        )\n        self.arc_face = ArcMarginProduct(CFG['embedding_size'], n_classes)\n\n    def forward(self, x, label=None):\n        emb = self.backbone(x)\n        if label is not None:\n            return self.arc_face(emb, label)\n        return F.normalize(emb, p=2, dim=1)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ====================================================\n# Run\n# ====================================================\ndef run_pipeline():\n    print(\"Loading CSV files...\")\n    train_df = pd.read_csv(CFG['train_csv'])\n    test_pairs = pd.read_csv(CFG['test_csv'])\n    \n    le = LabelEncoder()\n    train_df['label'] = le.fit_transform(train_df['ground_truth'])\n    n_classes = len(le.classes_)\n    print(f\"Number of identities found: {n_classes}\")\n\n    transform_train = T.Compose([\n        T.Resize((CFG['img_size'], CFG['img_size'])),\n        T.RandomHorizontalFlip(),\n        T.ColorJitter(brightness=0.2, contrast=0.2),\n        T.ToTensor(),\n        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n    ])\n\n    transform_test = T.Compose([\n        T.Resize((CFG['img_size'], CFG['img_size'])),\n        T.ToTensor(),\n        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n    ])\n\n    train_dataset = JaguarDataset(train_df, CFG['train_dir'], transform_train)\n    train_loader = DataLoader(train_dataset, batch_size=CFG['batch_size'], shuffle=True, num_workers=2)\n\n    model = JaguarModel(n_classes).to(CFG['device'])\n    optimizer = torch.optim.AdamW(\n        model.parameters(),\n        lr=CFG['lr'],\n        weight_decay=CFG['weight_decay']\n    )\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=CFG['epochs'])\n    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)\n\n    print(f\"Start Training: {CFG['model_name']} (Gradient Accumulation enabled)\")\n\n    # Accumulation step settings\n    accumulation_steps = CFG.get('accumulation_steps', 4)\n    \n    for epoch in range(CFG['epochs']):\n        model.train()\n        epoch_loss = 0\n        optimizer.zero_grad()\n        pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{CFG['epochs']}\")\n        \n        for i, (imgs, labels) in enumerate(pbar):\n            imgs, labels = imgs.to(CFG['device']), labels.to(CFG['device'])\n            # Forward\n            outputs = model(imgs, labels)\n            loss = criterion(outputs, labels)\n            # Average loss by the number of abbumulation steps\n            loss = loss / accumulation_steps\n            loss.backward()\n            if (i + 1) % accumulation_steps == 0:\n                optimizer.step()\n                optimizer.zero_grad()\n            \n            epoch_loss += loss.item() * accumulation_steps\n            pbar.set_postfix({'loss': epoch_loss / (i + 1)})\n\n        if (len(train_loader) % accumulation_steps) != 0:\n            optimizer.step()\n            optimizer.zero_grad()\n        \n        scheduler.step()\n\n    print(\"Extracting test embeddings...\")\n    col1 = 'query_image'\n    col2 = 'gallery_image'\n\n    unique_test_files = pd.concat([test_pairs[col1], test_pairs[col2]]).unique()\n    test_df_unique = pd.DataFrame({'filename': unique_test_files})\n    test_dataset = JaguarDataset(test_df_unique, CFG['test_dir'], transform_test, is_test=True)\n    test_loader = DataLoader(test_dataset, batch_size=CFG['batch_size'], shuffle=False)\n\n    model.eval()\n    emb_dict = {}\n    \n    scales = [448, 518] \n    \n    with torch.no_grad():\n        for imgs, fnames in tqdm(test_loader, desc=\"Multi-Scale TTA\"):\n            imgs = imgs.to(CFG['device'])\n            \n            multi_scale_embs = []\n            \n            for size in scales:\n                if size == CFG['img_size']:\n                    input_imgs = imgs\n                else:\n                    input_imgs = F.interpolate(imgs, size=(size, size), mode='bilinear', align_corners=False)\n                \n                # Normal + Horizontal Flip\n                emb_orig = model(input_imgs)\n                emb_flip = model(torch.flip(input_imgs, dims=[3]))\n                \n                multi_scale_embs.append(emb_orig)\n                multi_scale_embs.append(emb_flip)\n            \n            # Average all embeddings and normalize\n            emb_final = torch.stack(multi_scale_embs).mean(0)\n            emb_final = F.normalize(emb_final, p=2, dim=1)\n            \n            embs = emb_final.cpu().numpy()\n            for fname, emb in zip(fnames, embs):\n                emb_dict[fname] = emb\n\n    # --- Calculating similarity for submission ---\n    print(\"Calculating similarity for submission...\")\n    similarities = []\n    for index, row in tqdm(test_pairs.iterrows(), total=len(test_pairs), desc=\"Calculating\"):\n        emb1 = emb_dict[row[col1]]\n        emb2 = emb_dict[row[col2]]\n        # Dot product\n        score = np.dot(emb1, emb2)\n        score = (score + 1) / 2\n        similarities.append(score)\n    \n    # Submission\n    submission = pd.read_csv(CFG['sample_sub'])\n    submission['similarity'] = similarities\n    submission.to_csv('submission.csv', index=False)\n    \n    print(\"---------------------------------------\")\n    print(\"Success: 'submission.csv' created.\")\n    print(submission.head())\n\nif __name__ == \"__main__\":\n    run_pipeline()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}