{"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":"import warnings\nwarnings.filterwarnings(\"ignore\")\n\n!pip install wildlife-datasets git+https://github.com/WildlifeDatasets/wildlife-tools --quiet --upgrade-strategy only-if-needed\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torchvision.transforms as T\nimport timm\nimport optuna\nimport gc\nfrom transformers import AutoModel\nfrom sklearn.cluster import DBSCAN\nfrom sklearn.preprocessing import normalize\nfrom wildlife_datasets.datasets import AnimalCLEF2026\nfrom wildlife_tools.features import DeepFeatures\n\noptuna.logging.set_verbosity(optuna.logging.WARNING)\n\n# =============================================================================\n# SETTINGS\n# =============================================================================\nROOT_DIR   = \"/kaggle/input/animal-clef-2026\"\nDEVICE     = \"cuda\"\nBATCH_SIZE = 32\nN_TRIALS   = 40   # Optuna trials per species — raise for better tuning\n\n# Which model(s) to use per species + Optuna search range for eps\n# eps_start = official baseline value (seeded as first Optuna trial)\nSPECIES_CONFIG = {\n    \"LynxID2025\": {\n        \"models\": [\n            {\"name\": \"miewid\",         \"size\": 512},\n            {\"name\": \"megadescriptor\", \"size\": 384},\n        ],\n        \"eps_range\": (0.15, 0.50),\n        \"eps_start\": 0.30,\n    },\n    \"SalamanderID2025\": {\n        \"models\": [\n            {\"name\": \"megadescriptor\", \"size\": 384},\n            {\"name\": \"miewid\",         \"size\": 512},\n        ],\n        \"eps_range\": (0.10, 0.40),\n        \"eps_start\": 0.20,\n    },\n    \"SeaTurtleID2022\": {\n        \"models\": [\n            {\"name\": \"megadescriptor\", \"size\": 384},\n            {\"name\": \"miewid\",         \"size\": 512},\n        ],\n        \"eps_range\": (0.25, 0.60),\n        \"eps_start\": 0.40,\n    },\n    \"TexasHornedLizards\": {\n        \"models\": [\n            {\"name\": \"miewid\",         \"size\": 512},\n            {\"name\": \"megadescriptor\", \"size\": 384},\n        ],\n        \"eps_range\": (0.10, 0.45),\n        \"eps_start\": 0.24,\n    },\n}\n\n# =============================================================================\n# MODEL LOADING\n# =============================================================================\n\ndef load_model(name: str):\n    print(f\"    Loading {name}...\")\n    if name == \"miewid\":\n        model = AutoModel.from_pretrained(\n            \"conservationxlabs/miewid-msv3\", trust_remote_code=True\n        )\n    else:\n        model = timm.create_model(\n            \"hf-hub:BVRA/MegaDescriptor-L-384\", pretrained=True\n        )\n        model.reset_classifier(0)\n    return model.eval()\n\n\n# =============================================================================\n# FEATURE EXTRACTION\n# =============================================================================\n\ndef extract_features(model, dataset, size: int) -> np.ndarray:\n    \"\"\"Extract L2-normalized embeddings using wildlife_tools DeepFeatures.\"\"\"\n    transform = T.Compose([\n        T.Resize((size, size)),\n        T.ToTensor(),\n        T.Normalize(mean=(0.485, 0.456, 0.406),\n                    std=(0.229, 0.224, 0.225)),\n    ])\n    dataset.set_transform(transform)\n    extractor = DeepFeatures(model=model, device=DEVICE, batch_size=BATCH_SIZE)\n    result    = extractor(dataset)\n    feats     = result.features   # (N, D) numpy array\n    return normalize(feats, axis=1, norm=\"l2\")\n\n\n# =============================================================================\n# ENSEMBLE  (concatenate scaled embeddings — handles different dims)\n# =============================================================================\n\ndef ensemble_concat(feature_list: list, weights: list) -> np.ndarray:\n    \"\"\"\n    Scale each embedding block by its weight then concatenate column-wise.\n    Safe when models have different output dims (e.g. MiewID=2152, Mega=1536).\n    Final vector is L2-normalized.\n    \"\"\"\n    w      = np.array(weights, dtype=np.float32)\n    w     /= w.sum()\n    scaled = [wi * fi for wi, fi in zip(w, feature_list)]\n    return normalize(np.concatenate(scaled, axis=1), axis=1, norm=\"l2\")\n\n\n# =============================================================================\n# QUERY EXPANSION\n# =============================================================================\n\ndef query_expansion(features: np.ndarray, k: int, alpha: float) -> np.ndarray:\n    \"\"\"Average Query Expansion: blend each embedding with its k-NN mean.\"\"\"\n    sim = np.dot(features, features.T)\n    knn = np.argsort(sim, axis=1)[:, -k:]\n    out = np.zeros_like(features)\n    for i in range(len(features)):\n        out[i] = features[i] + alpha * features[knn[i]].mean(axis=0)\n    return normalize(out, axis=1, norm=\"l2\")\n\n\n# =============================================================================\n# DBSCAN CLUSTERING  (same approach as official starter notebook)\n# =============================================================================\n\ndef relabel_negatives(labels: np.ndarray) -> np.ndarray:\n    \"\"\"Give each DBSCAN noise point (-1) its own unique cluster id.\"\"\"\n    labels  = np.array(labels)\n    neg_idx = np.where(labels == -1)[0]\n    new_ids = np.arange(labels.max() + 1, labels.max() + 1 + len(neg_idx))\n    labels[neg_idx] = new_ids\n    return labels\n\n\ndef run_dbscan(features: np.ndarray, eps: float) -> np.ndarray:\n    \"\"\"Compute cosine distance matrix then run DBSCAN.\"\"\"\n    sim      = np.dot(features, features.T)\n    distance = (np.max(sim) - np.maximum(sim, 0)) / np.max(sim)\n    labels   = DBSCAN(eps=eps, metric=\"precomputed\", min_samples=2).fit_predict(distance)\n    return relabel_negatives(labels)\n\n\n# =============================================================================\n# OPTUNA OBJECTIVE\n# =============================================================================\n\ndef make_objective(feature_list: list, n_models: int, eps_range: tuple):\n    \"\"\"\n    Tunes: eps, qe_k, qe_alpha, per-model weights.\n    Score = fraction of clusters that contain >1 image (rewards meaningful groups).\n    \"\"\"\n    n = feature_list[0].shape[0]\n\n    def objective(trial: optuna.Trial) -> float:\n        eps     = trial.suggest_float(\"eps\",      *eps_range)\n        qe_k    = trial.suggest_int(\"qe_k\",       3, min(15, n - 1))\n        qe_a    = trial.suggest_float(\"qe_alpha\", 0.1, 0.8)\n        weights = [trial.suggest_float(f\"w{i}\", 0.05, 1.0)\n                   for i in range(n_models)]\n\n        feats      = ensemble_concat(feature_list, weights)\n        feats      = query_expansion(feats, k=qe_k, alpha=qe_a)\n        labels     = run_dbscan(feats, eps)\n        n_clusters = len(np.unique(labels))\n\n        # Prune useless extremes\n        if n_clusters <= 1 or n_clusters >= n:\n            raise optuna.exceptions.TrialPruned()\n\n        # Reward clusters that contain more than one image\n        counts         = np.bincount(labels)\n        multi_clusters = np.sum(counts > 1)\n        score          = multi_clusters / n_clusters\n\n        # Penalise if cluster count is unreasonably extreme\n        if n_clusters < n // 10:\n            score *= 0.5\n\n        return score\n\n    return objective\n\n\n# =============================================================================\n# PER-SPECIES PIPELINE\n# =============================================================================\n\ndef process_species(species: str, cfg: dict, dataset) -> pd.DataFrame:\n    n          = len(dataset)\n    model_cfgs = cfg[\"models\"]\n    n_models   = len(model_cfgs)\n\n    print(f\"\\n{'='*60}\")\n    print(f\"  {species}  ({n} images)\")\n    print(f\"{'='*60}\")\n\n    # ── 1. Extract features ──────────────────────────────────────────────────\n    feature_list = []\n    for mcfg in model_cfgs:\n        model = load_model(mcfg[\"name\"])\n        feats = extract_features(model, dataset, mcfg[\"size\"])\n        feature_list.append(feats)\n        print(f\"    {mcfg['name']} features: {feats.shape}\")\n        del model\n        torch.cuda.empty_cache()\n        gc.collect()\n\n    # ── 2. Optuna tuning ─────────────────────────────────────────────────────\n    print(f\"\\n  Tuning with Optuna ({N_TRIALS} trials)...\")\n\n    study = optuna.create_study(\n        direction=\"maximize\",\n        sampler=optuna.samplers.TPESampler(seed=42),\n        pruner=optuna.pruners.MedianPruner(n_startup_trials=8),\n    )\n    # Seed first trial at the official baseline so we always have a good result\n    # even if subsequent trials don't improve things\n    study.enqueue_trial({\n        \"eps\":      cfg[\"eps_start\"],\n        \"qe_k\":     5,\n        \"qe_alpha\": 0.5,\n        **{f\"w{i}\": 1.0 for i in range(n_models)},\n    })\n    study.optimize(\n        make_objective(feature_list, n_models, cfg[\"eps_range\"]),\n        n_trials=N_TRIALS,\n        show_progress_bar=False,\n    )\n\n    best = study.best_params\n    print(f\"  Best params : {best}\")\n    print(f\"  Best score  : {study.best_value:.4f}\")\n\n    # ── 3. Final prediction ──────────────────────────────────────────────────\n    best_weights = [best[f\"w{i}\"] for i in range(n_models)]\n    final_feats  = ensemble_concat(feature_list, best_weights)\n    final_feats  = query_expansion(\n        final_feats, k=best[\"qe_k\"], alpha=best[\"qe_alpha\"])\n    labels     = run_dbscan(final_feats, best[\"eps\"])\n    n_clusters = len(np.unique(labels))\n    print(f\"  Clusters found: {n_clusters}\")\n\n    # Cluster names MUST follow this format for the competition scorer\n    cluster_names = [f\"cluster_{species}_{lbl}\" for lbl in labels]\n    image_ids     = dataset.metadata[\"image_id\"].values\n\n    return pd.DataFrame({\"image_id\": image_ids, \"cluster\": cluster_names})\n\n\n# =============================================================================\n# MAIN\n# =============================================================================\n\ndef main():\n    print(\"📂  Loading AnimalCLEF2026...\")\n    full_dataset = AnimalCLEF2026(\n        ROOT_DIR,\n        transform=None,\n        load_label=True,\n        factorize_label=True,\n        check_files=False,\n    )\n\n    # Use .df for filtering (as shown in official starter notebook)\n    test_dataset = full_dataset.get_subset(full_dataset.df[\"split\"] == \"test\")\n\n    print(f\"Metadata columns : {list(test_dataset.metadata.columns)}\")\n    print(f\"Species in test  : {test_dataset.metadata['dataset'].unique()}\")\n    print(f\"Total test images: {len(test_dataset)}\")\n\n    all_results = []\n    for species in test_dataset.metadata[\"dataset\"].unique():\n        cfg = SPECIES_CONFIG.get(species)\n        if cfg is None:\n            print(f\"\\n⚠️  No config for '{species}' — skipping.\")\n            continue\n\n        sp_dataset = test_dataset.get_subset(\n            test_dataset.df[\"dataset\"] == species\n        )\n        result = process_species(species, cfg, sp_dataset)\n        all_results.append(result)\n\n    # ── Submission ───────────────────────────────────────────────────────────\n    preds      = pd.concat(all_results, ignore_index=True)\n    sample_sub = pd.read_csv(f\"{ROOT_DIR}/sample_submission.csv\")\n\n    pred_map = dict(zip(preds[\"image_id\"], preds[\"cluster\"]))\n    sample_sub[\"cluster\"] = (\n        sample_sub[\"image_id\"]\n        .map(pred_map)\n        .fillna(\"cluster_unknown_0\")\n    )\n\n    n_missing = sample_sub[\"cluster\"].isna().sum()\n    if n_missing:\n        print(f\"\\n⚠️  WARNING: {n_missing} rows missing cluster labels!\")\n    else:\n        print(f\"\\n✅  All {len(sample_sub)} rows have cluster labels.\")\n\n    out = \"/kaggle/working/submission.csv\"\n    sample_sub.to_csv(out, index=False)\n    print(f\"🏆  Saved: {out}\")\n    print(sample_sub.head(10))\n    print(\"\\nSample cluster name counts (top 10):\")\n    print(sample_sub[\"cluster\"].value_counts().head(10))\n\n\nmain()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-09T17:58:07.414788Z","iopub.execute_input":"2026-04-09T17:58:07.415577Z","iopub.status.idle":"2026-04-09T18:05:28.811422Z","shell.execute_reply.started":"2026-04-09T17:58:07.415543Z","shell.execute_reply":"2026-04-09T18:05:28.810407Z"}},"outputs":[],"execution_count":null}]}