{"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":31259,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# =============================================================================\n# JAGUAR RE-ID V3: DINOv2 + ArcFace (Single Cell - No Ordering Issues)\n# =============================================================================\nimport os, random, warnings\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset\nfrom PIL import Image\nimport torchvision.transforms as T\nimport timm\nfrom tqdm import tqdm\nfrom sklearn.model_selection import StratifiedShuffleSplit\nwarnings.filterwarnings('ignore')\n\nprint(f\"PyTorch: {torch.__version__}\")\nprint(f\"CUDA: {torch.cuda.is_available()}\")\nif torch.cuda.is_available():\n    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n\n# =============================================================================\n# CONFIG\n# =============================================================================\nclass CFG:\n    INPUT_DIR = \"/kaggle/input/jaguar-re-id\"\n    TRAIN_DIR = os.path.join(INPUT_DIR, \"train/train\")\n    TEST_DIR = os.path.join(INPUT_DIR, \"test/test\")\n    MODEL_NAME = 'vit_small_patch14_dinov2.lvd142m'\n    IMG_SIZE = 518\n    BATCH_SIZE = 16\n    EPOCHS = 12\n    LR = 1e-5\n    WEIGHT_DECAY = 1e-4\n    ARCFACE_SCALE = 30\n    ARCFACE_MARGIN = 0.5\n    VAL_RATIO = 0.0\n    DEVICE = torch.device('cuda')\n    SEED = 42\n\nrandom.seed(CFG.SEED)\nnp.random.seed(CFG.SEED)\ntorch.manual_seed(CFG.SEED)\ntorch.cuda.manual_seed_all(CFG.SEED)\nprint(f\"Model: {CFG.MODEL_NAME} @ {CFG.IMG_SIZE}px\")\n\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 = img_dir\n        self.transform = transform\n        self.is_test = is_test\n    def __len__(self):\n        return len(self.df)\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        image = Image.open(os.path.join(self.img_dir, row['filename'])).convert('RGB')\n        if self.transform:\n            image = self.transform(image)\n        if self.is_test:\n            return image, row['filename']\n        return image, torch.tensor(row['label'], dtype=torch.long)\n\ntrain_transform = T.Compose([\n    T.Resize((CFG.IMG_SIZE, CFG.IMG_SIZE)),\n    T.RandomHorizontalFlip(p=0.5),\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])\nval_transform = 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# =============================================================================\n# ARCFACE\n# =============================================================================\nclass ArcMarginProduct(nn.Module):\n    def __init__(self, in_features, out_features, s=30.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, self.m = s, m\n        self.cos_m, self.sin_m = np.cos(m), np.sin(m)\n        self.th, self.mm = np.cos(np.pi - m), np.sin(np.pi - m) * m\n    def forward(self, input, label):\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight)).clamp(-1+1e-7, 1-1e-7)\n        sine = torch.sqrt(1.0 - cosine ** 2)\n        phi = cosine * self.cos_m - sine * self.sin_m\n        phi = torch.where(cosine > self.th, phi, cosine - self.mm)\n        one_hot = torch.zeros(cosine.size(), device=input.device)\n        one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n        return ((one_hot * phi) + ((1.0 - one_hot) * cosine)) * self.s\n\n# =============================================================================\n# MODEL\n# =============================================================================\nclass JaguarModel(nn.Module):\n    def __init__(self, model_name, n_classes, s=30.0, m=0.50):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=True, num_classes=0, dynamic_img_size=True)\n        self.embedding_dim = self.backbone.num_features\n        self.arcface = ArcMarginProduct(self.embedding_dim, n_classes, s=s, m=m)\n        print(f\"Embedding dim: {self.embedding_dim}\")\n    def forward(self, x, label=None):\n        emb = self.backbone(x)\n        if label is not None:\n            return self.arcface(emb, label)\n        return F.normalize(emb, p=2, dim=1)\n    def get_embedding(self, x):\n        return F.normalize(self.backbone(x), p=2, dim=1)\n\n# =============================================================================\n# LOAD DATA\n# =============================================================================\nprint(\"\\nLoading data...\")\ntrain_df = pd.read_csv(os.path.join(CFG.INPUT_DIR, \"train.csv\"))\ntest_df = pd.read_csv(os.path.join(CFG.INPUT_DIR, \"test.csv\"))\nsample_sub = pd.read_csv(os.path.join(CFG.INPUT_DIR, \"sample_submission.csv\"))\n\nlabel_to_idx = {l: i for i, l in enumerate(sorted(train_df['ground_truth'].unique()))}\ntrain_df['label'] = train_df['ground_truth'].map(label_to_idx)\nn_classes = len(label_to_idx)\nprint(f\"Images: {len(train_df)}, Classes: {n_classes}, Test pairs: {len(test_df):,}\")\n\n# =============================================================================\n# DATALOADER\n# =============================================================================\ntrain_dataset = JaguarDataset(train_df, CFG.TRAIN_DIR, train_transform)\ntrain_loader = DataLoader(train_dataset, batch_size=CFG.BATCH_SIZE, shuffle=True, num_workers=2)\nprint(f\"Train batches: {len(train_loader)}\")\n\n# =============================================================================\n# INITIALIZE MODEL\n# =============================================================================\nprint(f\"\\nInitializing {CFG.MODEL_NAME}...\")\nmodel = JaguarModel(CFG.MODEL_NAME, n_classes, s=CFG.ARCFACE_SCALE, m=CFG.ARCFACE_MARGIN).to(CFG.DEVICE)\noptimizer = torch.optim.AdamW(model.parameters(), lr=CFG.LR, weight_decay=CFG.WEIGHT_DECAY)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=CFG.EPOCHS)\ncriterion = nn.CrossEntropyLoss()\nn_params = sum(p.numel() for p in model.parameters())\nprint(f\"Parameters: {n_params:,}\")\n\n# =============================================================================\n# TRAINING\n# =============================================================================\nprint(\"\\n\" + \"=\"*50 + \"\\nTRAINING\\n\" + \"=\"*50)\nfor epoch in range(CFG.EPOCHS):\n    model.train()\n    total_loss = 0\n    pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{CFG.EPOCHS}\")\n    for imgs, labels in pbar:\n        imgs, labels = imgs.to(CFG.DEVICE), labels.to(CFG.DEVICE)\n        optimizer.zero_grad()\n        loss = criterion(model(imgs, labels), labels)\n        loss.backward()\n        optimizer.step()\n        total_loss += loss.item()\n        pbar.set_postfix({'loss': total_loss / (pbar.n + 1)})\n    scheduler.step()\n    avg_loss = total_loss / len(train_loader)\n    print(f\"Epoch {epoch+1}/{CFG.EPOCHS} - Loss: {avg_loss:.4f}\")\n\nprint(\"\\nTraining complete!\")\n\n# =============================================================================\n# INFERENCE WITH TTA\n# =============================================================================\nprint(\"\\n\" + \"=\"*50 + \"\\nINFERENCE\\n\" + \"=\"*50)\ntest_images = sorted(set(test_df['query_image']) | set(test_df['gallery_image']))\nprint(f\"Unique test images: {len(test_images)}\")\n\ntest_data = pd.DataFrame({'filename': test_images})\ntest_dataset = JaguarDataset(test_data, CFG.TEST_DIR, val_transform, is_test=True)\ntest_loader = DataLoader(test_dataset, batch_size=CFG.BATCH_SIZE, shuffle=False, num_workers=2)\n\nemb_dict = {}\nmodel.eval()\nwith torch.no_grad():\n    for imgs, fnames in tqdm(test_loader, desc=\"Extracting embeddings\"):\n        imgs = imgs.to(CFG.DEVICE)\n        emb_orig = model.get_embedding(imgs)\n        emb_flip = model.get_embedding(torch.flip(imgs, dims=[3]))\n        emb = F.normalize(emb_orig + emb_flip, p=2, dim=1)\n        for fname, e in zip(fnames, emb.cpu().numpy()):\n            emb_dict[fname] = e\n\nprint(f\"Extracted {len(emb_dict)} embeddings\")\n\n# =============================================================================\n# GENERATE SUBMISSION\n# =============================================================================\nprint(\"\\nGenerating submission...\")\nsimilarities = []\nfor _, row in tqdm(test_df.iterrows(), total=len(test_df), desc=\"Computing similarities\"):\n    emb1 = emb_dict[row['query_image']]\n    emb2 = emb_dict[row['gallery_image']]\n    sim = (np.dot(emb1, emb2) + 1) / 2\n    similarities.append(sim)\n\nsubmission = sample_sub.copy()\nsubmission['similarity'] = similarities\nsubmission.to_csv('submission.csv', index=False)\n\nprint(\"\\n\" + \"=\"*50)\nprint(\"DONE!\")\nprint(\"=\"*50)\nprint(f\"Similarity range: [{submission['similarity'].min():.4f}, {submission['similarity'].max():.4f}]\")\nprint(f\"Mean: {submission['similarity'].mean():.4f}, Std: {submission['similarity'].std():.4f}\")\nprint(f\"\\nSubmission saved to submission.csv\")\nprint(submission.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T21:17:01.839058Z","iopub.execute_input":"2026-02-01T21:17:01.839477Z","iopub.status.idle":"2026-02-01T22:34:51.205554Z","shell.execute_reply.started":"2026-02-01T21:17:01.839422Z","shell.execute_reply":"2026-02-01T22:34:51.204764Z"}},"outputs":[],"execution_count":null}]}