{"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":"gpu","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 -q --upgrade wandb\n\nimport os\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport timm\nfrom torchvision import transforms\nfrom PIL import Image\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.metrics.pairwise import cosine_similarity\nfrom tqdm.notebook import tqdm\nimport wandb\nimport math\nimport random\n\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nos.environ[\"HF_TOKEN\"] = user_secrets.get_secret(\"hf_api\")\nos.environ[\"WANDB_API_KEY\"] = user_secrets.get_secret(\"wandb_api\")\n\n# Fix all random seeds\nSEED = 42\ntorch.manual_seed(SEED)\nnp.random.seed(SEED)\nrandom.seed(SEED)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Device: {device}\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-03-02T21:33:04.324603Z","iopub.execute_input":"2026-03-02T21:33:04.325118Z","iopub.status.idle":"2026-03-02T21:33:33.270975Z","shell.execute_reply.started":"2026-03-02T21:33:04.325091Z","shell.execute_reply":"2026-03-02T21:33:33.270170Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"config = {\n    # Paths\n    \"data_dir\": Path(\"/kaggle/input/competitions/jaguar-re-id\"),\n    \"checkpoint_dir\": Path(\"checkpoints\"),\n    \n    # Model\n    \"megadescriptor_model\": \"hf-hub:BVRA/MegaDescriptor-L-384\",\n    \"input_size\": 384,\n    \"embedding_dim\": 256,\n    \"hidden_dim\": 512,\n    \n    # Training — FIXED across all loss functions\n    \"batch_size\": 32,\n    \"learning_rate\": 1e-4,\n    \"weight_decay\": 1e-4,\n    \"num_epochs\": 25,  # reduced from 50 to save GPU time\n    \"patience\": 7,\n    \"val_split\": 0.2,\n    \"seed\": SEED,\n    \n    # ArcFace defaults (used as reference)\n    \"arcface_margin\": 0.5,\n    \"arcface_scale\": 64.0,\n    \"dropout\": 0.3,\n}\n\nconfig[\"checkpoint_dir\"].mkdir(exist_ok=True)\nprint(\"Config ready\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T21:35:34.950142Z","iopub.execute_input":"2026-03-02T21:35:34.950831Z","iopub.status.idle":"2026-03-02T21:35:34.956115Z","shell.execute_reply.started":"2026-03-02T21:35:34.950799Z","shell.execute_reply":"2026-03-02T21:35:34.955381Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load training data\ntrain_df = pd.read_csv(config[\"data_dir\"] / \"train.csv\")\n\nprint(f\"Total images: {len(train_df)}\")\nprint(f\"Unique identities: {train_df['ground_truth'].nunique()}\")\n\n# Encode labels\nlabel_encoder = LabelEncoder()\ntrain_df['label_encoded'] = label_encoder.fit_transform(train_df['ground_truth'])\nnum_classes = len(label_encoder.classes_)\n\n# Stratified split — same seed every time, this never changes\ntrain_data, val_data = train_test_split(\n    train_df,\n    test_size=config[\"val_split\"],\n    random_state=config[\"seed\"],\n    stratify=train_df['ground_truth']\n)\n\nprint(f\"Train: {len(train_data)} images\")\nprint(f\"Val: {len(val_data)} images\")\nprint(f\"Num classes: {num_classes}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T21:35:40.248315Z","iopub.execute_input":"2026-03-02T21:35:40.248637Z","iopub.status.idle":"2026-03-02T21:35:40.281209Z","shell.execute_reply.started":"2026-03-02T21:35:40.248612Z","shell.execute_reply":"2026-03-02T21:35:40.280646Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load MegaDescriptor\nprint(\"Loading MegaDescriptor...\")\nmegadescriptor = timm.create_model(\n    config[\"megadescriptor_model\"],\n    pretrained=True\n)\nmegadescriptor.eval()\nmegadescriptor.to(device)\n\n# Get embedding dimension\nwith torch.no_grad():\n    dummy = torch.randn(1, 3, config[\"input_size\"], config[\"input_size\"]).to(device)\n    megadescriptor_dim = megadescriptor(dummy).shape[1]\n\nprint(f\"MegaDescriptor embedding dim: {megadescriptor_dim}\")\n\n# Preprocessing\npreprocess = transforms.Compose([\n    transforms.Resize((config[\"input_size\"], config[\"input_size\"])),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                         std=[0.229, 0.224, 0.225]),\n])\n\n@torch.no_grad()\ndef extract_embeddings(model, image_paths, batch_size=32):\n    model.eval()\n    embeddings = []\n    for i in tqdm(range(0, len(image_paths), batch_size)):\n        batch_paths = image_paths[i:i+batch_size]\n        tensors = []\n        for p in batch_paths:\n            try:\n                img = Image.open(p).convert(\"RGB\")\n                tensors.append(preprocess(img))\n            except:\n                tensors.append(torch.zeros(3, config[\"input_size\"], config[\"input_size\"]))\n        batch = torch.stack(tensors).to(device)\n        embeddings.append(model(batch).cpu().numpy())\n    return np.vstack(embeddings)\n\n# Extract and cache\ncache_dir = Path(\"embeddings\")\ncache_dir.mkdir(exist_ok=True)\n\ntrain_paths = [config[\"data_dir\"] / \"train/train\" / f for f in train_data[\"filename\"].values]\nval_paths = [config[\"data_dir\"] / \"train/train\" / f for f in val_data[\"filename\"].values]\n\ntrain_cache = cache_dir / \"train_embeddings.npz\"\nval_cache = cache_dir / \"val_embeddings.npz\"\n\nif train_cache.exists():\n    train_embeddings = np.load(train_cache)[\"embeddings\"]\n    print(f\"Loaded cached train embeddings: {train_embeddings.shape}\")\nelse:\n    print(\"Extracting train embeddings...\")\n    train_embeddings = extract_embeddings(megadescriptor, train_paths)\n    np.savez_compressed(train_cache, embeddings=train_embeddings)\n    print(f\"Saved train embeddings: {train_embeddings.shape}\")\n\nif val_cache.exists():\n    val_embeddings = np.load(val_cache)[\"embeddings\"]\n    print(f\"Loaded cached val embeddings: {val_embeddings.shape}\")\nelse:\n    print(\"Extracting val embeddings...\")\n    val_embeddings = extract_embeddings(megadescriptor, val_paths)\n    np.savez_compressed(val_cache, embeddings=val_embeddings)\n    print(f\"Saved val embeddings: {val_embeddings.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T21:36:27.572105Z","iopub.execute_input":"2026-03-02T21:36:27.572693Z","iopub.status.idle":"2026-03-02T21:48:37.610661Z","shell.execute_reply.started":"2026-03-02T21:36:27.572662Z","shell.execute_reply":"2026-03-02T21:48:37.609993Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class EmbeddingDataset(Dataset):\n    def __init__(self, embeddings, labels):\n        self.embeddings = torch.FloatTensor(embeddings)\n        self.labels = torch.LongTensor(labels)\n    \n    def __len__(self):\n        return len(self.labels)\n    \n    def __getitem__(self, idx):\n        return self.embeddings[idx], self.labels[idx]\n\n\ndef compute_validation_map(model, val_embeddings, val_labels):\n    model.eval()\n    with torch.no_grad():\n        val_tensor = torch.FloatTensor(val_embeddings).to(device)\n        finetuned_emb = model.get_embeddings(val_tensor).cpu().numpy()\n    \n    sim_matrix = cosine_similarity(finetuned_emb)\n    np.fill_diagonal(sim_matrix, -1)\n    \n    identity_aps = {}\n    for query_idx in range(len(val_labels)):\n        query_label = val_labels[query_idx]\n        similarities = sim_matrix[query_idx]\n        is_match = (val_labels == query_label).astype(int)\n        is_match[query_idx] = 0\n        \n        sorted_indices = np.argsort(-similarities)\n        sorted_matches = is_match[sorted_indices]\n        \n        n_positives = sorted_matches.sum()\n        if n_positives == 0:\n            continue\n        \n        cumsum = np.cumsum(sorted_matches)\n        precision_at_k = cumsum / np.arange(1, len(sorted_matches) + 1)\n        ap = np.sum(precision_at_k * sorted_matches) / n_positives\n        \n        if query_label not in identity_aps:\n            identity_aps[query_label] = []\n        identity_aps[query_label].append(ap)\n    \n    identity_mean_aps = [np.mean(aps) for aps in identity_aps.values()]\n    return np.mean(identity_mean_aps)\n\nprint(\"Dataset and validation function ready\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T21:49:31.469827Z","iopub.execute_input":"2026-03-02T21:49:31.470558Z","iopub.status.idle":"2026-03-02T21:49:31.478654Z","shell.execute_reply.started":"2026-03-02T21:49:31.470527Z","shell.execute_reply":"2026-03-02T21:49:31.477905Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ── 1. ArcFace (baseline loss) ──────────────────────────────────────────\nclass ArcFaceLayer(nn.Module):\n    def __init__(self, embedding_dim, num_classes, margin=0.5, scale=64.0):\n        super().__init__()\n        self.margin = margin\n        self.scale = scale\n        self.weight = nn.Parameter(torch.FloatTensor(num_classes, embedding_dim))\n        nn.init.xavier_uniform_(self.weight)\n\n    def forward(self, embeddings, labels):\n        embeddings = F.normalize(embeddings, dim=1)\n        weight = F.normalize(self.weight, dim=1)\n        cosine = F.linear(embeddings, weight).clamp(-1+1e-6, 1-1e-6)\n        theta = torch.acos(cosine)\n        one_hot = torch.zeros_like(cosine)\n        one_hot.scatter_(1, labels.view(-1, 1), 1.0)\n        output = self.scale * torch.cos(theta + one_hot * self.margin)\n        return F.cross_entropy(output, labels)\n\n\n# ── 2. CosFace ───────────────────────────────────────────────────────────\nclass CosFaceLayer(nn.Module):\n    def __init__(self, embedding_dim, num_classes, margin=0.35, scale=64.0):\n        super().__init__()\n        self.margin = margin\n        self.scale = scale\n        self.weight = nn.Parameter(torch.FloatTensor(num_classes, embedding_dim))\n        nn.init.xavier_uniform_(self.weight)\n\n    def forward(self, embeddings, labels):\n        embeddings = F.normalize(embeddings, dim=1)\n        weight = F.normalize(self.weight, dim=1)\n        cosine = F.linear(embeddings, weight)\n        one_hot = torch.zeros_like(cosine)\n        one_hot.scatter_(1, labels.view(-1, 1), 1.0)\n        output = self.scale * (cosine - one_hot * self.margin)\n        return F.cross_entropy(output, labels)\n\n\n# ── 3. Triplet Loss (hard mining) ────────────────────────────────────────\nclass TripletLoss(nn.Module):\n    def __init__(self, margin=0.3):\n        super().__init__()\n        self.margin = margin\n\n    def forward(self, embeddings, labels):\n        embeddings = F.normalize(embeddings, dim=1)\n        dist = torch.cdist(embeddings, embeddings, p=2)\n        pos_mask = labels.unsqueeze(0) == labels.unsqueeze(1)\n        neg_mask = ~pos_mask\n        pos_mask.fill_diagonal_(False)\n\n        hardest_pos = (dist * pos_mask.float()).max(dim=1)[0]\n        hardest_neg = (dist + 1e6 * (~neg_mask).float()).min(dim=1)[0]\n\n        loss = F.relu(hardest_pos - hardest_neg + self.margin)\n        return loss.mean()\n\n\n# ── 4. ArcFace + Triplet Combined ────────────────────────────────────────\nclass CombinedLoss(nn.Module):\n    def __init__(self, embedding_dim, num_classes, alpha=0.5):\n        super().__init__()\n        self.arcface = ArcFaceLayer(embedding_dim, num_classes)\n        self.triplet = TripletLoss()\n        self.alpha = alpha\n\n    def forward(self, embeddings, labels):\n        return (self.alpha * self.arcface(embeddings, labels) +\n                (1 - self.alpha) * self.triplet(embeddings, labels))\n\n\n# ── 5. Focal Loss ────────────────────────────────────────────────────────\nclass FocalLoss(nn.Module):\n    def __init__(self, embedding_dim, num_classes, gamma=2.0, scale=64.0):\n        super().__init__()\n        self.gamma = gamma\n        self.scale = scale\n        self.weight = nn.Parameter(torch.FloatTensor(num_classes, embedding_dim))\n        nn.init.xavier_uniform_(self.weight)\n\n    def forward(self, embeddings, labels):\n        embeddings = F.normalize(embeddings, dim=1)\n        weight = F.normalize(self.weight, dim=1)\n        logits = self.scale * F.linear(embeddings, weight)\n        ce = F.cross_entropy(logits, labels, reduction='none')\n        pt = torch.exp(-ce)\n        loss = ((1 - pt) ** self.gamma * ce).mean()\n        return loss\n\n\n# ── 6. SubCenter ArcFace ─────────────────────────────────────────────────\nclass SubCenterArcFace(nn.Module):\n    def __init__(self, embedding_dim, num_classes, K=3, margin=0.5, scale=64.0):\n        super().__init__()\n        self.K = K\n        self.num_classes = num_classes\n        self.margin = margin\n        self.scale = scale\n        self.weight = nn.Parameter(\n            torch.FloatTensor(num_classes * K, embedding_dim))\n        nn.init.xavier_uniform_(self.weight)\n\n    def forward(self, embeddings, labels):\n        embeddings = F.normalize(embeddings, dim=1)\n        weight = F.normalize(self.weight, dim=1)\n        cosine = F.linear(embeddings, weight)\n        cosine = cosine.view(-1, self.num_classes, self.K).max(dim=2)[0]\n        cosine = cosine.clamp(-1+1e-6, 1-1e-6)\n        theta = torch.acos(cosine)\n        one_hot = torch.zeros_like(cosine)\n        one_hot.scatter_(1, labels.view(-1, 1), 1.0)\n        output = self.scale * torch.cos(theta + one_hot * self.margin)\n        return F.cross_entropy(output, labels)\n\nprint(\"All 6 loss functions defined\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T21:51:02.096480Z","iopub.execute_input":"2026-03-02T21:51:02.097052Z","iopub.status.idle":"2026-03-02T21:51:02.114532Z","shell.execute_reply.started":"2026-03-02T21:51:02.097023Z","shell.execute_reply":"2026-03-02T21:51:02.113808Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class EmbeddingProjection(nn.Module):\n    def __init__(self, input_dim=1536, hidden_dim=512, \n                 output_dim=256, dropout=0.3):\n        super().__init__()\n        self.network = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim),\n            nn.ReLU(inplace=True),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, output_dim),\n            nn.BatchNorm1d(output_dim),\n        )\n    \n    def forward(self, x):\n        return self.network(x)\n    \n    def get_embeddings(self, x):\n        return F.normalize(self.forward(x), p=2, dim=1)\n\nprint(\"Embedding model defined\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T21:51:53.963409Z","iopub.execute_input":"2026-03-02T21:51:53.964004Z","iopub.status.idle":"2026-03-02T21:51:53.969860Z","shell.execute_reply.started":"2026-03-02T21:51:53.963976Z","shell.execute_reply":"2026-03-02T21:51:53.969174Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_one_epoch(model, loss_fn, loader, optimizer, device):\n    model.train()\n    loss_fn.train()\n    total_loss = 0\n\n    for embeddings, labels in tqdm(loader, desc=\"Training\", leave=False):\n        embeddings, labels = embeddings.to(device), labels.to(device)\n\n        optimizer.zero_grad()\n        projected = model(embeddings)\n        loss = loss_fn(projected, labels)\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n\n    return total_loss / len(loader)\n\n\ndef validate_one_epoch(model, loss_fn, loader, device):\n    model.eval()\n    loss_fn.eval()\n    total_loss = 0\n\n    with torch.no_grad():\n        for embeddings, labels in tqdm(loader, desc=\"Validation\", leave=False):\n            embeddings, labels = embeddings.to(device), labels.to(device)\n            projected = model(embeddings)\n            loss = loss_fn(projected, labels)\n            total_loss += loss.item()\n\n    return total_loss / len(loader)\n\n\nprint(\"Training functions defined\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T21:52:27.058579Z","iopub.execute_input":"2026-03-02T21:52:27.059270Z","iopub.status.idle":"2026-03-02T21:52:27.065775Z","shell.execute_reply.started":"2026-03-02T21:52:27.059244Z","shell.execute_reply":"2026-03-02T21:52:27.065045Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define all 6 experiments\nexperiments = [\n    {\n        \"name\": \"loss-arcface\",\n        \"loss_fn\": lambda: ArcFaceLayer(config[\"embedding_dim\"], num_classes)\n    },\n    {\n        \"name\": \"loss-cosface\",\n        \"loss_fn\": lambda: CosFaceLayer(config[\"embedding_dim\"], num_classes)\n    },\n    {\n        \"name\": \"loss-triplet\",\n        \"loss_fn\": lambda: TripletLoss()\n    },\n    {\n        \"name\": \"loss-combined-arcface-triplet\",\n        \"loss_fn\": lambda: CombinedLoss(config[\"embedding_dim\"], num_classes)\n    },\n    {\n        \"name\": \"loss-focal\",\n        \"loss_fn\": lambda: FocalLoss(config[\"embedding_dim\"], num_classes)\n    },\n    {\n        \"name\": \"loss-subcenter-arcface\",\n        \"loss_fn\": lambda: SubCenterArcFace(config[\"embedding_dim\"], num_classes)\n    },\n]\n\n# Fixed datasets — same for all experiments\ntrain_dataset = EmbeddingDataset(train_embeddings, train_data['label_encoded'].values)\nval_dataset = EmbeddingDataset(val_embeddings, val_data['label_encoded'].values)\n\ntrain_loader = DataLoader(train_dataset, batch_size=config[\"batch_size\"], \n                          shuffle=True, num_workers=0)\nval_loader = DataLoader(val_dataset, batch_size=config[\"batch_size\"], \n                        shuffle=False, num_workers=0)\n\nval_labels = val_data['ground_truth'].values\n\n# Results table\nresults = []\n\n# ── Main Loop ────────────────────────────────────────────────────────────\nfor exp in experiments:\n    print(f\"\\n{'='*60}\")\n    print(f\"Running: {exp['name']}\")\n    print(f\"{'='*60}\")\n\n    # Reset seeds before each experiment for fairness\n    torch.manual_seed(SEED)\n    np.random.seed(SEED)\n    random.seed(SEED)\n\n    # Initialize model and loss fresh for each experiment\n    model = EmbeddingProjection(\n        input_dim=megadescriptor_dim,\n        hidden_dim=config[\"hidden_dim\"],\n        output_dim=config[\"embedding_dim\"],\n        dropout=config[\"dropout\"]\n    ).to(device)\n\n    loss_fn = exp[\"loss_fn\"]().to(device)\n\n    # Combine model and loss parameters for optimizer\n    all_params = list(model.parameters()) + list(loss_fn.parameters())\n    optimizer = torch.optim.AdamW(\n        all_params,\n        lr=config[\"learning_rate\"],\n        weight_decay=config[\"weight_decay\"]\n    )\n\n    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n        optimizer, mode='min', factor=0.5, patience=3\n    )\n\n    # Initialize W&B run\n    wandb.login(key=os.environ[\"WANDB_API_KEY\"])\n    run = wandb.init(\n        project=\"jaguar-reid-mishank\",\n        name=exp[\"name\"],\n        config={\n            **config,\n            \"loss_function\": exp[\"name\"],\n            \"num_parameters\": sum(p.numel() for p in model.parameters()),\n            \"num_classes\": num_classes,\n        }\n    )\n\n    # Training loop\n    best_val_loss = float('inf')\n    best_val_map = 0.0\n    patience_counter = 0\n\n    for epoch in range(config[\"num_epochs\"]):\n        train_loss = train_one_epoch(model, loss_fn, train_loader, optimizer, device)\n        val_loss = validate_one_epoch(model, loss_fn, val_loader, device)\n        val_map = compute_validation_map(model, val_embeddings, val_labels)\n\n        scheduler.step(val_loss)\n        current_lr = optimizer.param_groups[0]['lr']\n\n        # Log to W&B\n        wandb.log({\n            \"epoch\": epoch + 1,\n            \"train_loss\": train_loss,\n            \"val_loss\": val_loss,\n            \"val_map\": val_map,\n            \"learning_rate\": current_lr,\n        })\n\n        print(f\"Epoch {epoch+1:2d} | \"\n              f\"Train Loss: {train_loss:.4f} | \"\n              f\"Val Loss: {val_loss:.4f} | \"\n              f\"Val mAP: {val_map:.4f}\")\n\n        # Save best model\n        if val_loss < best_val_loss:\n            best_val_loss = val_loss\n            best_val_map = val_map\n            patience_counter = 0\n            checkpoint_path = config[\"checkpoint_dir\"] / f\"{exp['name']}_best.pth\"\n            torch.save({\n                \"model_state_dict\": model.state_dict(),\n                \"val_map\": best_val_map,\n                \"val_loss\": best_val_loss,\n            }, checkpoint_path)\n        else:\n            patience_counter += 1\n            if patience_counter >= config[\"patience\"]:\n                print(f\"Early stopping at epoch {epoch+1}\")\n                break\n\n    # Log best results and save artifact\n    wandb.log({\n        \"best_val_map\": best_val_map,\n        \"best_val_loss\": best_val_loss,\n    })\n\n    artifact = wandb.Artifact(f\"model-{exp['name']}\", type=\"model\")\n    artifact.add_file(str(checkpoint_path))\n    wandb.log_artifact(artifact)\n\n    wandb.finish()\n\n    # Store results\n    results.append({\n        \"loss_function\": exp[\"name\"],\n        \"best_val_map\": best_val_map,\n        \"best_val_loss\": best_val_loss,\n    })\n\n    print(f\"Finished {exp['name']} | Best Val mAP: {best_val_map:.4f}\")\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"ALL EXPERIMENTS COMPLETE\")\nprint(\"=\"*60)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T21:56:13.126807Z","iopub.execute_input":"2026-03-02T21:56:13.127421Z","iopub.status.idle":"2026-03-02T21:58:26.250646Z","shell.execute_reply.started":"2026-03-02T21:56:13.127382Z","shell.execute_reply":"2026-03-02T21:58:26.249831Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"results_df = pd.DataFrame(results)\nresults_df = results_df.sort_values(\"best_val_map\", ascending=False)\nresults_df[\"best_val_map\"] = results_df[\"best_val_map\"].round(4)\nresults_df[\"best_val_loss\"] = results_df[\"best_val_loss\"].round(4)\n\nprint(\"Loss Function Comparison Results\")\nprint(\"=\"*50)\nprint(results_df.to_string(index=False))\nprint(\"=\"*50)\nprint(f\"Best loss function: {results_df.iloc[0]['loss_function']}\")\nprint(f\"Best val mAP: {results_df.iloc[0]['best_val_map']}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T21:59:01.540168Z","iopub.execute_input":"2026-03-02T21:59:01.540848Z","iopub.status.idle":"2026-03-02T21:59:01.555979Z","shell.execute_reply.started":"2026-03-02T21:59:01.540816Z","shell.execute_reply":"2026-03-02T21:59:01.555431Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load best model (combined loss)\nbest_model = EmbeddingProjection(\n    input_dim=megadescriptor_dim,\n    hidden_dim=config[\"hidden_dim\"],\n    output_dim=config[\"embedding_dim\"],\n    dropout=config[\"dropout\"]\n).to(device)\n\ncheckpoint = torch.load(\n    config[\"checkpoint_dir\"] / \"loss-combined-arcface-triplet_best.pth\",\n    map_location=device,\n    weights_only=False\n)\nbest_model.load_state_dict(checkpoint[\"model_state_dict\"])\nbest_model.eval()\nprint(f\"Loaded best model | Val mAP: {checkpoint['val_map']:.4f}\")\n\n# Load test pairs\ntest_pairs_df = pd.read_csv(config[\"data_dir\"] / \"test.csv\")\ntest_images = sorted(list(\n    set(test_pairs_df['query_image'].unique()) | \n    set(test_pairs_df['gallery_image'].unique())\n))\n\n# Extract test embeddings\ntest_paths = [config[\"data_dir\"] / \"test/test\" / f for f in test_images]\nprint(f\"Extracting embeddings for {len(test_images)} test images...\")\ntest_mega_embeddings = extract_embeddings(megadescriptor, test_paths)\n\nwith torch.no_grad():\n    test_tensor = torch.FloatTensor(test_mega_embeddings).to(device)\n    test_embeddings_proj = best_model.get_embeddings(test_tensor).cpu().numpy()\n\n# Compute similarities\nimg_to_emb = {fn: emb for fn, emb in zip(test_images, test_embeddings_proj)}\n\nsimilarities = []\nfor _, row in tqdm(test_pairs_df.iterrows(), total=len(test_pairs_df)):\n    q = img_to_emb[row['query_image']]\n    g = img_to_emb[row['gallery_image']]\n    similarities.append(float(np.dot(q, g)))\n\nsimilarities = np.clip(similarities, 0.0, 1.0)\n\nsubmission_df = pd.DataFrame({\n    'row_id': test_pairs_df['row_id'],\n    'similarity': similarities\n})\n\nsubmission_df.to_csv(\"submission_combined_loss.csv\", index=False)\nprint(f\"Submission saved: {len(submission_df)} rows\")\n\nfrom IPython.display import FileLink\nFileLink('submission_combined_loss.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T22:01:29.104047Z","iopub.execute_input":"2026-03-02T22:01:29.104811Z","iopub.status.idle":"2026-03-02T22:03:57.348712Z","shell.execute_reply.started":"2026-03-02T22:01:29.104779Z","shell.execute_reply":"2026-03-02T22:03:57.348028Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}