{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":127282,"databundleVersionId":15520105,"sourceType":"competition"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install wildlife-datasets wildlife-tools timm scikit-learn --quiet --upgrade","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Imports\n","metadata":{}},{"cell_type":"code","source":"import os\nimport warnings\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torchvision.transforms as T\nfrom torch.utils.data import DataLoader\nfrom tqdm.notebook import tqdm\nfrom sklearn.cluster import AgglomerativeClustering\nfrom sklearn.preprocessing import normalize\nimport timm\nfrom transformers import AutoModel\nfrom wildlife_datasets.datasets import AnimalCLEF2026","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Config\n","metadata":{}},{"cell_type":"code","source":"warnings.filterwarnings(\"ignore\")\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nBATCH_SIZE = 64\nNUM_WORKERS = 4\nROOT_DIR = \"/kaggle/input/animal-clef-2026\"\n\nSPECIES_CONFIG = {\n    \"LynxID2025\": {\n        \"model_type\": \"miewid\",\n        \"threshold\": 0.35,\n        \"image_size\": 512,\n        \"qe_k\": 5,\n    },\n    \"SalamanderID2025\": {\n        \"model_type\": \"megadescriptor\",\n        \"threshold\": 0.25,\n        \"image_size\": 384,\n        \"qe_k\": 3,\n    },\n    \"SeaTurtleID2022\": {\n        \"model_type\": \"megadescriptor\",\n        \"threshold\": 0.45,\n        \"image_size\": 384,\n        \"qe_k\": 8,\n    },\n    \"TexasHornedLizards\": {\n        \"model_type\": \"miewid\",\n        \"threshold\": 0.30,\n        \"image_size\": 512,\n        \"qe_k\": 5,\n    },\n}\n\nif torch.cuda.device_count() > 1:\n    print(f\"🚀 Using {torch.cuda.device_count()} GPUs\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Models\n","metadata":{}},{"cell_type":"code","source":"class MiewIDWrapper(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.backbone = AutoModel.from_pretrained(\n            \"conservationxlabs/miewid-msv3\", trust_remote_code=True\n        )\n\n    def forward(self, x):\n        return self.backbone(x)\n\n\nclass MegaDescriptorWrapper(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.backbone = timm.create_model(\n            \"hf-hub:BVRA/MegaDescriptor-L-384\", pretrained=True\n        )\n        self.backbone.reset_classifier(0)\n\n    def forward(self, x):\n        return self.backbone(x)\n\n\ndef get_model(model_type):\n    if model_type == \"miewid\":\n        model = MiewIDWrapper()\n    else:\n        model = MegaDescriptorWrapper()\n\n    model.to(DEVICE)\n    if torch.cuda.device_count() > 1:\n        model = nn.DataParallel(model)\n    model.eval()\n    return model","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Features\n","metadata":{}},{"cell_type":"code","source":"def extract_features(model, dataset, image_size):\n    dataset.transform = T.Compose(\n        [\n            T.Resize((image_size, image_size)),\n            T.ToTensor(),\n            T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ]\n    )\n    loader = DataLoader(\n        dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, shuffle=False\n    )\n    all_features = []\n\n    with torch.no_grad():\n        for batch in tqdm(loader, leave=False):\n            images = batch[0].to(DEVICE)\n            with torch.cuda.amp.autocast():\n                feats_sum = model(images) + model(torch.flip(images, dims=[3]))\n\n            feats_norm = torch.nn.functional.normalize(feats_sum.float(), p=2, dim=1)\n            all_features.append(feats_norm.cpu().numpy())\n\n    return np.concatenate(all_features)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Processing\n","metadata":{}},{"cell_type":"code","source":"def query_expansion(features, k=5, alpha=0.5):\n    sim_matrix = np.dot(features, features.T)\n    knn_indices = np.argsort(sim_matrix, axis=1)[:, -k:]\n    expanded = np.zeros_like(features)\n    for i in range(features.shape[0]):\n        expanded[i] = features[i] + alpha * np.mean(features[knn_indices[i]], axis=0)\n    return normalize(expanded, axis=1, norm=\"l2\")\n\n\ndef cluster_and_format(features, image_ids, species, config):\n    print(f\"    -> Clustering {species} (Thresh: {config['threshold']})...\")\n    dist_matrix = np.clip(1.0 - np.dot(features, features.T), 0, 1)\n\n    labels = AgglomerativeClustering(\n        n_clusters=None,\n        metric=\"precomputed\",\n        linkage=\"average\",\n        distance_threshold=config[\"threshold\"],\n    ).fit_predict(dist_matrix)\n\n    print(f\"    ✅ Found {len(np.unique(labels))} individuals.\")\n    return pd.DataFrame(\n        {\n            \"image_id\": image_ids,\n            \"cluster\": [f\"cluster_{species}_{lbl}\" for lbl in labels],\n        }\n    )","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Execution\n","metadata":{}},{"cell_type":"code","source":"print(\"📂 Loading Data...\")\nfull_dataset = AnimalCLEF2026(root=ROOT_DIR, transform=None, load_label=True)\ntest_meta = full_dataset.metadata[full_dataset.metadata[\"split\"] == \"test\"]\nresults = []\n\nfor species in test_meta[\"dataset\"].unique():\n    print(f\"    🦁 Processing: {species}\")\n    cfg = SPECIES_CONFIG.get(species)\n\n    # Subset\n    sp_meta = test_meta[test_meta[\"dataset\"] == species]\n    sp_dataset = full_dataset.get_subset(sp_meta.index.values)\n\n    # Process\n    model = get_model(cfg[\"model_type\"])\n    feats = extract_features(model, sp_dataset, cfg[\"image_size\"])\n    feats = query_expansion(feats, k=cfg[\"qe_k\"])\n\n    # Store\n    results.append(cluster_and_format(feats, sp_meta[\"image_id\"].values, species, cfg))\n\n    del model, feats, sp_dataset\n    torch.cuda.empty_cache()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Submission\n","metadata":{}},{"cell_type":"code","source":"preds = pd.concat(results)\nsample_sub = pd.read_csv(os.path.join(ROOT_DIR, \"sample_submission.csv\"))\n\n# Map predictions\npred_map = dict(zip(preds[\"image_id\"], preds[\"cluster\"]))\nsample_sub[\"cluster\"] = sample_sub[\"image_id\"].map(pred_map).fillna(\"cluster_error_0\")\n\nsample_sub.to_csv(\"submission.csv\", index=False)\nprint(\"🏆 'submission.csv' generated!\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pd.read_csv(\"/kaggle/working/submission.csv\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}