{"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":[{"sourceType":"competition","sourceId":127282,"databundleVersionId":15520105}],"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,"execution":{"iopub.status.busy":"2026-03-28T06:51:36.498633Z","iopub.execute_input":"2026-03-28T06:51:36.498812Z","iopub.status.idle":"2026-03-28T06:51:51.982409Z","shell.execute_reply.started":"2026-03-28T06:51:36.498793Z","shell.execute_reply":"2026-03-28T06:51:51.981715Z"}},"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,"execution":{"iopub.status.busy":"2026-03-28T06:51:51.983487Z","iopub.execute_input":"2026-03-28T06:51:51.984435Z","iopub.status.idle":"2026-03-28T06:52:17.166235Z","shell.execute_reply.started":"2026-03-28T06:51:51.984401Z","shell.execute_reply":"2026-03-28T06:52:17.165602Z"}},"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,"execution":{"iopub.status.busy":"2026-03-28T06:52:17.168229Z","iopub.execute_input":"2026-03-28T06:52:17.168820Z","iopub.status.idle":"2026-03-28T06:52:17.453523Z","shell.execute_reply.started":"2026-03-28T06:52:17.168791Z","shell.execute_reply":"2026-03-28T06:52:17.452860Z"}},"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,"execution":{"iopub.status.busy":"2026-03-28T06:52:17.456842Z","iopub.execute_input":"2026-03-28T06:52:17.457098Z","iopub.status.idle":"2026-03-28T06:52:17.463010Z","shell.execute_reply.started":"2026-03-28T06:52:17.457080Z","shell.execute_reply":"2026-03-28T06:52:17.462139Z"}},"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,"execution":{"iopub.status.busy":"2026-03-28T06:52:17.463846Z","iopub.execute_input":"2026-03-28T06:52:17.464091Z","iopub.status.idle":"2026-03-28T06:52:17.485669Z","shell.execute_reply.started":"2026-03-28T06:52:17.464072Z","shell.execute_reply":"2026-03-28T06:52:17.485166Z"}},"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,"execution":{"iopub.status.busy":"2026-03-28T06:52:17.486667Z","iopub.execute_input":"2026-03-28T06:52:17.487160Z","iopub.status.idle":"2026-03-28T06:52:17.500905Z","shell.execute_reply.started":"2026-03-28T06:52:17.487131Z","shell.execute_reply":"2026-03-28T06:52:17.500416Z"}},"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,"execution":{"iopub.status.busy":"2026-03-28T06:52:17.501616Z","iopub.execute_input":"2026-03-28T06:52:17.501824Z","iopub.status.idle":"2026-03-28T06:57:09.300021Z","shell.execute_reply.started":"2026-03-28T06:52:17.501807Z","shell.execute_reply":"2026-03-28T06:57:09.299258Z"}},"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,"execution":{"iopub.status.busy":"2026-03-28T06:57:09.302202Z","iopub.execute_input":"2026-03-28T06:57:09.303255Z","iopub.status.idle":"2026-03-28T06:57:09.337133Z","shell.execute_reply.started":"2026-03-28T06:57:09.303218Z","shell.execute_reply":"2026-03-28T06:57:09.336558Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pd.read_csv(\"/kaggle/working/submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-28T06:57:09.337849Z","iopub.execute_input":"2026-03-28T06:57:09.338117Z","iopub.status.idle":"2026-03-28T06:57:09.366569Z","shell.execute_reply.started":"2026-03-28T06:57:09.338096Z","shell.execute_reply":"2026-03-28T06:57:09.366005Z"}},"outputs":[],"execution_count":null}]}