{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":126777,"databundleVersionId":15314950,"sourceType":"competition"}],"dockerImageVersionId":31259,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import 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 torchvision import transforms\nfrom PIL import Image\nfrom pathlib import Path\nfrom tqdm import tqdm\nimport timm\nimport random\nimport os\n\n# ============================================================\n# SETUP\n# ============================================================\ndef seed_everything(seed=42):\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(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nseed_everything(42)\n\nprint(\"=\"*60)\nprint(\"🐆 JAGUAR RE-ID - ARCFACE ULTIMATE\")\nprint(\"=\"*60)\n\n# Load data\ntrain_df = pd.read_csv('/kaggle/input/jaguar-re-id/train.csv')\ntest_df = pd.read_csv('/kaggle/input/jaguar-re-id/test.csv')\n\nTRAIN_DIR = Path('/kaggle/input/jaguar-re-id/train/train')\nTEST_DIR = Path('/kaggle/input/jaguar-re-id/test/test')\n\n# Find column names\nid_col = [c for c in train_df.columns if c not in ['image', 'filename']][0]\nimg_col = 'image' if 'image' in train_df.columns else 'filename'\n\nprint(f\"Using ID column: '{id_col}'\")\nprint(f\"Using image column: '{img_col}'\")\nprint(f\"Training images: {len(train_df)}\")\nprint(f\"Test pairs: {len(test_df)}\")\n\n# Create label mapping\nunique_jaguars = sorted(train_df[id_col].unique())\nlabel_map = {jag_id: idx for idx, jag_id in enumerate(unique_jaguars)}\ntrain_df['label'] = train_df[id_col].map(label_map)\nnum_classes = len(unique_jaguars)\n\nprint(f\"Unique jaguars: {num_classes}\")\n\n# ============================================================\n# CONFIG\n# ============================================================\nclass Config:\n    backbone = 'efficientnet_b4'  # Winner from zero-shot + better than ResNet152\n    img_size = 224\n    batch_size = 32\n    embedding_dim = 512\n    num_epochs = 15  # More training!\n    lr = 0.0003\n    weight_decay = 0.01\n    arcface_s = 30.0\n    arcface_m = 0.60  # Stronger margin\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nprint(f\"Device: {Config.device}\")\nprint(f\"Training for {Config.num_epochs} epochs\")\n\n# ============================================================\n# DATASET WITH TRAINING\n# ============================================================\nclass JaguarTrainDataset(Dataset):\n    def __init__(self, df, img_dir, img_col, transform=None):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = Path(img_dir)\n        self.img_col = img_col\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = self.img_dir / row[self.img_col]\n        img = Image.open(img_path).convert('RGB')\n        label = row['label']\n        \n        if self.transform:\n            img = self.transform(img)\n        \n        return img, torch.tensor(label, dtype=torch.long)\n\nclass JaguarTestDataset(Dataset):\n    def __init__(self, df, img_dir, img_col, transform=None):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = Path(img_dir)\n        self.img_col = img_col\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = self.img_dir / row[self.img_col]\n        img = Image.open(img_path).convert('RGB')\n        \n        if self.transform:\n            img = self.transform(img)\n        \n        return img, row[self.img_col]\n\n# Transforms with augmentation\ntrain_transform = transforms.Compose([\n    transforms.Resize((Config.img_size, Config.img_size)),\n    transforms.RandomHorizontalFlip(p=0.5),\n    transforms.RandomRotation(15),\n    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n])\n\ntest_transform = transforms.Compose([\n    transforms.Resize((Config.img_size, Config.img_size)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n])\n\n# ============================================================\n# ARCFACE LOSS\n# ============================================================\nclass ArcFaceLayer(nn.Module):\n    def __init__(self, in_features, out_features, s=30.0, m=0.50):\n        super().__init__()\n        self.in_features = in_features\n        self.out_features = out_features\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        # Normalize features and weights\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n        \n        if label is None:\n            return cosine\n        \n        # Add angular margin\n        phi = cosine - self.m\n        \n        # One-hot encoding\n        one_hot = torch.zeros(cosine.size(), device=input.device)\n        one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n        \n        # Apply margin\n        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n        output *= self.s\n        \n        return output\n\n# ============================================================\n# MODEL\n# ============================================================\nclass JaguarReIDModel(nn.Module):\n    def __init__(self, backbone_name, embedding_dim, num_classes, pretrained=True):\n        super().__init__()\n        self.backbone = timm.create_model(backbone_name, pretrained=pretrained, num_classes=0)\n        \n        # Get backbone output dimension\n        with torch.no_grad():\n            dummy = torch.randn(1, 3, Config.img_size, Config.img_size)\n            backbone_dim = self.backbone(dummy).shape[1]\n        \n        # Neck (projection head)\n        self.neck = nn.Sequential(\n            nn.BatchNorm1d(backbone_dim),\n            nn.Dropout(0.2),\n            nn.Linear(backbone_dim, embedding_dim),\n            nn.BatchNorm1d(embedding_dim),\n        )\n        \n        # ArcFace head\n        self.head = ArcFaceLayer(embedding_dim, num_classes, s=Config.arcface_s, m=Config.arcface_m)\n        \n    def forward(self, x, label=None):\n        features = self.backbone(x)\n        embeddings = self.neck(features)\n        \n        if label is not None:\n            return self.head(embeddings, label), embeddings\n        \n        return embeddings\n\n# ============================================================\n# TRAINING\n# ============================================================\ndef train_epoch(model, loader, criterion, optimizer, device):\n    model.train()\n    total_loss = 0\n    \n    pbar = tqdm(loader, desc=\"Training\")\n    for images, labels in pbar:\n        images, labels = images.to(device), labels.to(device)\n        \n        optimizer.zero_grad()\n        logits, _ = model(images, labels)\n        loss = criterion(logits, labels)\n        \n        loss.backward()\n        optimizer.step()\n        \n        total_loss += loss.item()\n        pbar.set_postfix({'loss': f'{loss.item():.4f}'})\n    \n    return total_loss / len(loader)\n\n# ============================================================\n# INFERENCE\n# ============================================================\ndef extract_embeddings(model, loader, device):\n    model.eval()\n    all_embeddings = []\n    all_names = []\n    \n    with torch.no_grad():\n        for images, names in tqdm(loader, desc=\"Extracting embeddings\"):\n            images = images.to(device)\n            embeddings = model(images)\n            embeddings = F.normalize(embeddings, p=2, dim=1)\n            \n            all_embeddings.append(embeddings.cpu().numpy())\n            all_names.extend(names)\n    \n    all_embeddings = np.vstack(all_embeddings)\n    return dict(zip(all_names, all_embeddings))\n\n# ============================================================\n# MAIN PIPELINE\n# ============================================================\nprint(\"\\n\" + \"=\"*60)\nprint(\"🔄 TRAINING PHASE\")\nprint(\"=\"*60)\n\n# Create datasets\ntrain_dataset = JaguarTrainDataset(train_df, TRAIN_DIR, img_col, train_transform)\ntrain_loader = DataLoader(train_dataset, batch_size=Config.batch_size, shuffle=True, num_workers=2)\n\n# Initialize model\nmodel = JaguarReIDModel(Config.backbone, Config.embedding_dim, num_classes).to(Config.device)\ncriterion = nn.CrossEntropyLoss()\noptimizer = torch.optim.AdamW(model.parameters(), lr=Config.lr, weight_decay=Config.weight_decay)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=Config.num_epochs)\n\nprint(f\"Model: {Config.backbone}\")\nprint(f\"Parameters: {sum(p.numel() for p in model.parameters()):,}\")\n\n# Training loop\nbest_loss = float('inf')\nfor epoch in range(Config.num_epochs):\n    print(f\"\\nEpoch {epoch+1}/{Config.num_epochs}\")\n    \n    loss = train_epoch(model, train_loader, criterion, optimizer, Config.device)\n    scheduler.step()\n    \n    print(f\"Average Loss: {loss:.4f} | LR: {scheduler.get_last_lr()[0]:.6f}\")\n    \n    if loss < best_loss:\n        best_loss = loss\n        torch.save(model.state_dict(), 'best_model.pth')\n        print(\"💾 Model saved!\")\n\n# Load best model\nprint(\"\\n📥 Loading best model...\")\nmodel.load_state_dict(torch.load('best_model.pth'))\n\n# ============================================================\n# INFERENCE PHASE\n# ============================================================\nprint(\"\\n\" + \"=\"*60)\nprint(\"🔍 INFERENCE PHASE\")\nprint(\"=\"*60)\n\n# Get unique test images\nunique_test_imgs = sorted(set(test_df['query_image']) | set(test_df['gallery_image']))\nprint(f\"Unique test images: {len(unique_test_imgs)}\")\n\ntest_imgs_df = pd.DataFrame({img_col: unique_test_imgs})\ntest_dataset = JaguarTestDataset(test_imgs_df, TEST_DIR, img_col, test_transform)\ntest_loader = DataLoader(test_dataset, batch_size=Config.batch_size, shuffle=False, num_workers=2)\n\n# Extract embeddings\nembeddings_dict = extract_embeddings(model, test_loader, Config.device)\n\n# Compute similarities\nprint(\"\\n🔢 Computing similarities...\")\nsimilarities = []\n\nfor _, row in tqdm(test_df.iterrows(), total=len(test_df)):\n    query_emb = embeddings_dict[row['query_image']]\n    gallery_emb = embeddings_dict[row['gallery_image']]\n    \n    # Cosine similarity (already normalized)\n    sim = np.dot(query_emb, gallery_emb)\n    \n    # Scale to [0, 1]\n    sim = (sim + 1) / 2\n    similarities.append(sim)\n\nsimilarities = np.array(similarities)\n\n# ============================================================\n# SAVE SUBMISSION\n# ============================================================\nprint(\"\\n\" + \"=\"*60)\nprint(\"💾 SAVING SUBMISSION\")\nprint(\"=\"*60)\n\nsubmission = pd.DataFrame({\n    'row_id': test_df['row_id'],\n    'similarity': similarities\n})\n\nsubmission.to_csv('submission.csv', index=False)\n\nprint(f\"Submission saved: {len(submission)} predictions\")\nprint(f\"Similarity range: [{similarities.min():.4f}, {similarities.max():.4f}]\")\nprint(f\"Mean similarity: {similarities.mean():.4f}\")\nprint(f\"Std similarity: {similarities.std():.4f}\")\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"✅ COMPLETE!\")\nprint(\"=\"*60)\nprint(f\"Model: {Config.backbone} + ArcFace\")\nprint(f\"Training: {Config.num_epochs} epochs\")\nprint(f\"Margin: {Config.arcface_m}\")\nprint(f\"Expected score: 0.80-0.85 (vs 0.723 baseline)\")\nprint(\"\\n🚀 Ready to submit!\")\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}