{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":28755,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Jaguar Re-Identification\n### ArcFace + MegaDescriptor\n\nRe-identifying individual jaguars from photos using **embeddings** (numerical fingerprints) compared with **cosine similarity**.\n\n**Approach:**\n- **Backbone:** MegaDescriptor (pre-trained on animal recognition)\n- **Fine-tuning:** ArcFace (angular margin to clearly separate identities)\n- **Imbalance:** weighted sampler + data augmentation\n- **Output:** similarity score [0-1] for each test image pair\n\n**Note:** training classifies the 31 train jaguars only as a *means* to build good embeddings; at test time the embeddings are used to compare jaguars never seen before.","metadata":{}},{"cell_type":"code","source":"!pip install timm -qq","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import timm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom tqdm import tqdm\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader, WeightedRandomSampler\nfrom PIL import Image\nfrom sklearn.metrics.pairwise import cosine_similarity\nfrom sklearn.model_selection import train_test_split\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport os","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if torch.cuda.is_available():\n    device = torch.device(\"cuda\")\n    print(\"Using GPU:\", torch.cuda.get_device_name(0))\nelse:\n    device = torch.device(\"cpu\")\n    print(\"WARNING: no GPU, using CPU (slow)\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"base = \"/kaggle/input/competitions/jaguar-re-id\"\n\nfor f in os.listdir(base):\n    print(f)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train = pd.read_csv(base + '/train.csv')\n\ndisplay(train.head())\ntrain.info()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train['label'] = train['ground_truth'].astype('category').cat.codes","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"photo_counts = train[\"ground_truth\"].value_counts()\n\nplt.figure(figsize=(14, 5))\nphoto_counts.plot(kind=\"bar\", color=\"steelblue\")\nplt.title(\"Photos per jaguar (class imbalance)\", fontsize=14)\nplt.xlabel(\"Jaguar\")\nplt.ylabel(\"Number of photos\")\nplt.xticks(rotation=75)\nplt.tight_layout()\nplt.show()\n\nprint(\"Most photographed:\", photo_counts.idxmax(), \"→\", photo_counts.max(), \"photos\")\nprint(\"Least photographed:\", photo_counts.idxmin(), \"→\", photo_counts.min(), \"photos\")\nprint(\"Ratio:\", round(photo_counts.max() / photo_counts.min(), 1), \"x\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"img_dir = base + '/train/train'\n# 6 foto dello stesso giaguaro\nname = \"Marcela\"\nphotos = train[train[\"ground_truth\"] == name][\"filename\"].values[:6]\n\nplt.figure(figsize=(15, 8))\nfor i, fname in enumerate(photos):\n    img = Image.open(f\"{img_dir}/{fname}\").convert(\"RGB\")\n    plt.subplot(2, 3, i + 1)\n    plt.imshow(img)\n    plt.title(f\"{name}\\n{fname}\")\n    plt.axis(\"off\")\nplt.suptitle(f\"Same jaguar ({name}), different photos\", fontsize=14)\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_split, val_split = train_test_split(\n    train,\n    test_size=0.2,\n    stratify=train['label'],\n    random_state=42\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"tmp_model = timm.create_model(\"hf-hub:BVRA/MegaDescriptor-L-384\", pretrained=True, num_classes=0)\ndata_config = timm.data.resolve_model_data_config(tmp_model)\n\ntransform_train = timm.data.create_transform(\n    **data_config, is_training=True, hflip=0.5, color_jitter=0.3\n)\ntransform_val = timm.data.create_transform(**data_config, is_training=False)\n\ndel tmp_model","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class JaguarDataset(Dataset):\n    def __init__(self, dataframe, img_dir, transform):\n        self.df = dataframe.reset_index(drop=True)\n        self.img_dir = img_dir\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 = Image.open(f\"{self.img_dir}/{row['filename']}\").convert(\"RGB\")\n        img = self.transform(img)\n        label = row[\"label\"]\n        return img, label","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_dataset = JaguarDataset(train_split, img_dir, transform_train)\nval_dataset = JaguarDataset(val_split, img_dir, transform_val)\n\nprint(\"Train:\", len(train_dataset), \"| Val:\", len(val_dataset))\n\n# Prova: la stessa foto presa due volte deve risultare leggermente diversa con augmentation\nimg1, _ = train_dataset[0]\nimg2, _ = train_dataset[0]\ndiff = (img1 - img2).abs().sum().item()\nprint(\"Difference between two draws of the same train image:\", round(diff, 2))\nprint(\"(if > 0, random augmentation is working)\")\n\nimg, label = train_dataset[0]\nprint(\"Prepared image shape:\", img.shape)\nprint(\"Label (jaguar number):\", label)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class_counts = np.bincount(train_split['label'].values)\nclass_weights = 1.0 / class_counts\nsample_weights = class_weights[train_split['label'].values]\n\nsampler = WeightedRandomSampler(\n    weights=sample_weights,\n    num_samples=len(sample_weights),\n    replacement=True\n)\n\nbatch_size = 32\n\ntrain_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=sampler, num_workers=2)\nval_loader   = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ArcFaceHead(nn.Module):\n    def __init__(self, embedding_dim, n_classes, margin=0.5, scale=30.0):\n        super().__init__()\n        self.margin = margin\n        self.scale = scale\n        self.weight = nn.Parameter(torch.empty(n_classes, embedding_dim))\n        nn.init.xavier_uniform_(self.weight)\n\n    def forward(self, embeddings, labels):\n        emb_norm = F.normalize(embeddings, dim=1)\n        w_norm = F.normalize(self.weight, dim=1)\n        cosine = emb_norm @ w_norm.t()\n        cosine = cosine.clamp(-1.0 + 1e-7, 1.0 - 1e-7)\n        angle = torch.acos(cosine)\n        margin_one_hot = F.one_hot(labels, num_classes=self.weight.shape[0]).float()\n        angle_with_margin = angle + margin_one_hot * self.margin\n        logits = torch.cos(angle_with_margin) * self.scale\n        return logits","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class JaguarModel(nn.Module):\n    def __init__(self, n_classes, embedding_dim=512):\n        super().__init__()\n        # MegaDescriptor\n        self.backbone = timm.create_model(\n            \"hf-hub:BVRA/MegaDescriptor-L-384\",\n            pretrained=True,\n            num_classes=0\n        )\n        backbone_dim = self.backbone.num_features\n\n        self.embedding = nn.Linear(backbone_dim, embedding_dim)\n\n        # ArcFace\n        self.arcface = ArcFaceHead(embedding_dim, n_classes)\n\n    def forward(self, images, labels):\n        features = self.backbone(images)\n        emb = self.embedding(features)\n        logits = self.arcface(emb, labels)\n        return logits\n\n    def get_embedding(self, images):\n        features = self.backbone(images)\n        emb = self.embedding(features)\n        return emb\n\nmodel = JaguarModel(n_classes=31).to(device)\n\n# Congelo MegaDescriptor tranne gli ultimi strati\nfor param in model.backbone.parameters():\n    param.requires_grad = False\nfor param in list(model.backbone.parameters())[-20:]:\n    param.requires_grad = True\n\ntrainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\ntotal = sum(p.numel() for p in model.parameters())\nprint(f\"Trainable params: {trainable:,}\")\nprint(f\"Total params:     {total:,}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"criterion = nn.CrossEntropyLoss()\n\noptimizer = optim.AdamW(\n    filter(lambda p: p.requires_grad, model.parameters()),\n    lr=1e-4\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"EPOCHS = 30\nPATIENCE = 3\nbest_val_loss = float(\"inf\")\npatience_counter = 0\n\nfor epoch in range(EPOCHS):\n\n    model.train()\n    train_loss_total = 0.0\n    for images, labels in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{EPOCHS} - train\"):\n        images = images.to(device)\n        labels = labels.to(device).long()\n        logits = model(images, labels)\n        loss = criterion(logits, labels)\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n        train_loss_total += loss.item()\n    mean_train_loss = train_loss_total / len(train_loader)\n\n    model.eval()\n    val_loss_total = 0.0\n    with torch.no_grad():\n        for images, labels in val_loader:\n            images = images.to(device)\n            labels = labels.to(device).long()\n            logits = model(images, labels)\n            loss = criterion(logits, labels)\n            val_loss_total += loss.item()\n    mean_val_loss = val_loss_total / len(val_loader)\n\n    print(f\"Epoch {epoch+1}: train_loss={mean_train_loss:.4f} | val_loss={mean_val_loss:.4f}\")\n\n    if mean_val_loss < best_val_loss:\n        best_val_loss = mean_val_loss\n        patience_counter = 0\n        torch.save(model.state_dict(), \"best_model.pth\")\n        print(\"  -> saved (best so far)\")\n    else:\n        patience_counter += 1\n        print(f\"  -> no improvement ({patience_counter}/{PATIENCE})\")\n        if patience_counter >= PATIENCE:\n            print(f\"\\nEarly stopping at epoch {epoch+1}: no improvement for {PATIENCE} epochs.\")\n            break","metadata":{"trusted":true,"execution":{"execution_failed":"2026-06-13T19:07:20.804Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.load_state_dict(torch.load(\"best_model.pth\"))\nmodel.eval()\nprint(\"Best model loaded.\")\n\ntest = pd.read_csv(f\"{base}/test.csv\")\nprint(\"Rows (pairs) in test:\", len(test))\nprint(\"Columns:\", list(test.columns))\nprint(test.head())","metadata":{"trusted":true,"execution":{"execution_failed":"2026-06-13T19:07:20.804Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_dir = f\"{base}/test/test\"\n\nunique_test_images = sorted(set(test[\"query_image\"]) | set(test[\"gallery_image\"]))\nprint(\"Unique test images:\", len(unique_test_images))\n\ntest_embeddings = {}\n\nmodel.eval()\nwith torch.no_grad():\n    for fname in tqdm(unique_test_images, desc=\"Extracting embeddings\"):\n        img = Image.open(f\"{test_dir}/{fname}\").convert(\"RGB\")\n        img = transform_val(img).unsqueeze(0).to(device)\n        emb = model.get_embedding(img)\n        emb = F.normalize(emb, dim=1)\n        test_embeddings[fname] = emb.cpu()\n\nprint(\"Embeddings computed:\", len(test_embeddings))\nprint(\"Embedding shape:\", test_embeddings[unique_test_images[0]].shape)","metadata":{"trusted":true,"execution":{"execution_failed":"2026-06-13T19:07:20.804Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"similarities = []\n\nfor _, row in tqdm(test.iterrows(), total=len(test), desc=\"Comparing pairs\"):\n    q_emb = test_embeddings[row[\"query_image\"]]\n    g_emb = test_embeddings[row[\"gallery_image\"]]\n\n    # Cosine similarity con impronte normalizzate\n    score = (q_emb * g_emb).sum().item()\n\n    # Da [-1, 1] a [0, 1]\n    score = (score + 1) / 2\n\n    similarities.append(score)\n\nsubmission = pd.DataFrame({\n    \"row_id\": test[\"row_id\"],\n    \"similarity\": similarities\n})\nsubmission[\"similarity\"] = submission[\"similarity\"].clip(0, 1)\nsubmission.to_csv(\"submission.csv\", index=False)","metadata":{"trusted":true,"execution":{"execution_failed":"2026-06-13T19:07:20.804Z"}},"outputs":[],"execution_count":null}]}