{"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":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":127282,"databundleVersionId":15520105}],"dockerImageVersionId":31328,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install git+https://github.com/WildlifeDatasets/wildlife-datasets@develop\n!pip install git+https://github.com/WildlifeDatasets/wildlife-tools\n!pip install hnswlib\n!pip install transformers==4.57.6","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2026-05-03T21:58:12.826083Z","iopub.execute_input":"2026-05-03T21:58:12.826368Z","iopub.status.idle":"2026-05-03T21:59:53.957046Z","shell.execute_reply.started":"2026-05-03T21:58:12.826335Z","shell.execute_reply":"2026-05-03T21:59:53.956067Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Dataset Loading","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\n\nimport os\n\nimport os\nimport numpy as np\nimport pandas as pd\nimport timm\nimport torch\nimport torchvision.transforms as T\nfrom wildlife_datasets.datasets import AnimalCLEF2026\nfrom wildlife_tools.features import DeepFeatures\nfrom wildlife_tools.similarity import CosineSimilarity\nfrom wildlife_tools.similarity.wildfusion import SimilarityPipeline, WildFusion\nfrom wildlife_tools.similarity.pairwise.lightglue import MatchLightGlue\nfrom wildlife_tools.features.local import AlikedExtractor, SuperPointExtractor, SiftExtractor, DiskExtractor\nfrom wildlife_tools.similarity.calibration import IsotonicCalibration\nfrom transformers import AutoModel\nfrom sklearn.model_selection import train_test_split\nimport hnswlib\nimport cv2\n\nimport warnings\nwarnings.filterwarnings('ignore')\n\nroot = '/kaggle/input/competitions/animal-clef-2026'\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T21:59:53.958224Z","iopub.execute_input":"2026-05-03T21:59:53.958501Z","iopub.status.idle":"2026-05-03T22:00:41.236545Z","shell.execute_reply.started":"2026-05-03T21:59:53.958473Z","shell.execute_reply":"2026-05-03T22:00:41.235342Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create a PyTorch dataset class that loads images, applies dataset-specific preprocessing, fixes orientation when metadata is available, and returns transformed images.\nfrom PIL import Image\nfrom torch.utils.data import Dataset\n\nclass WildlifeDataset(Dataset):\n    def __init__(self, df, transform=None):\n        self.metadata = df.reset_index(drop=True).copy()\n        self.transform = transform\n        self.col_label = 'identity'\n\n    def __len__(self):\n        return len(self.metadata)\n\n    def __getitem__(self, idx):\n        row = self.metadata.iloc[idx]\n        dataset_name = str(row.get('dataset', ''))\n    \n        img = Image.open(os.path.join(root, row[\"path\"])).convert(\"RGB\")\n    \n        # 🔹 Convert LynxID2025 to grayscale\n        if dataset_name == 'LynxID2025':\n            img = img.convert(\"L\").convert(\"RGB\")  # keep 3 channels\n    \n        # 1. Orientation Logic\n        orientation = str(row.get('orientation', 'nan'))\n    \n        if orientation != 'nan':\n            if orientation == 'left':\n                img = img.transpose(Image.FLIP_LEFT_RIGHT)\n    \n            elif orientation == 'top' and (dataset_name != 'SalamanderID2025' or img.size[1] > img.size[0] * 1.1):\n                img = img.transpose(Image.ROTATE_270)\n    \n            elif orientation == 'bottom' or (dataset_name == 'SalamanderID2025' and img.size[1] > img.size[0] * 1.1):\n                img = img.transpose(Image.ROTATE_90)\n            \n        # 2. Transforms\n        if self.transform is not None:\n            img = self.transform(img)\n    \n        identity = str(row[\"identity\"])\n    \n        return img, identity","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T22:00:41.246753Z","iopub.execute_input":"2026-05-03T22:00:41.247082Z","iopub.status.idle":"2026-05-03T22:00:41.320332Z","shell.execute_reply.started":"2026-05-03T22:00:41.247044Z","shell.execute_reply":"2026-05-03T22:00:41.319475Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Embedding Model Transformation\ntransform_miew = T.Compose([\n    T.Resize([512, 512]),\n    T.ToTensor(),\n    T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))\n])\n\n# Feature Matching Transformation\ntransforms_local = T.Compose([\n    T.Resize([512, 512]),\n    T.ToTensor()\n])","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Mixed Split to Simulate Wildlife Datasets","metadata":{}},{"cell_type":"code","source":"# Load metadata and build train, validation, eval, and unknown-ID splits.\n# The split keeps dataset balance and separates some identities as unknown.\ndef merge_rare_strat_keys(df, key_col, min_count=2):\n    counts = df[key_col].value_counts()\n    rare = counts[counts < min_count].index\n    df = df.copy()\n    df[key_col] = df[key_col].apply(lambda x: x.split('_')[0] + '_other' if x in rare else x)\n    \n    id_counts = df[key_col].value_counts().to_dict()\n    df[f'id_count_{key_col}'] = df[key_col].map(id_counts)\n    df = df[df[f\"id_count_{key_col}\"] > 1] \n    return df\n\ndf_metadata = pd.read_csv(os.path.join(root, 'metadata.csv'))\n\ndf_train = df_metadata[(df_metadata['split'] == 'train')].reset_index(drop=True)\ndf_test = df_metadata[(df_metadata['split'] == 'test')]\n\n## Known & Unknown ID Split\nid_counts = df_train['identity'].value_counts()\ndf_train['id_count'] = df_train['identity'].map(id_counts)\ndf_ids = df_train[['dataset','identity', 'id_count']].drop_duplicates().reset_index(drop=True)\ndf_ids['strat_key'] = df_ids['dataset'].astype(str)\ndf_ids = merge_rare_strat_keys(df_ids, 'strat_key')\n\ndf_train_tmp_ids, df_unknown_ids = train_test_split(\n    df_ids,\n    test_size=0.05,\n    random_state=42,\n    stratify=df_ids['strat_key']\n)\n\n\nunknown_ids_list = df_unknown_ids['identity'].unique().tolist()\ndf_unknown = df_train[df_train['identity'].isin(unknown_ids_list)]\n\n## Stratify Based on ID\ndf_train_tmp = df_train[~df_train['identity'].isin(unknown_ids_list)]\ndf_train_tmp['strat_key'] = df_train_tmp['identity'].astype(str)\ndf_train_tmp = merge_rare_strat_keys(df_train_tmp, 'strat_key')\n\ntrain_df, df_val_tmp = train_test_split(\n    df_train_tmp,\n    test_size=0.2,\n    random_state=42,\n    stratify=df_train_tmp['strat_key']\n)\n\nval_eval_df = pd.concat((df_val_tmp, df_unknown), axis=0, ignore_index=True)\n\nval_eval_df['strat_key'] = val_eval_df['dataset'].astype(str) + '_' + val_eval_df['identity'].astype(str)\nval_eval_df = merge_rare_strat_keys(val_eval_df, 'strat_key')\n\nval_df, eval_df = train_test_split(\n    val_eval_df,\n    test_size=0.5,\n    random_state=42,\n    stratify=val_eval_df['strat_key']\n)\nval_df['split'] = 'val'\neval_df['split'] = 'eval'\n\ndf_all_mixed_split = pd.concat((train_df, val_df, eval_df, df_test), axis=0, ignore_index=True)\n\ndataset_all = WildlifeDataset(df_all_mixed_split, transform=transform_miew)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T22:00:41.325085Z","iopub.execute_input":"2026-05-03T22:00:41.325423Z","iopub.status.idle":"2026-05-03T22:00:41.529447Z","shell.execute_reply.started":"2026-05-03T22:00:41.325400Z","shell.execute_reply":"2026-05-03T22:00:41.527750Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Embedding Generation & Retrieval","metadata":{}},{"cell_type":"code","source":"# Build an HNSW index from embeddings and retrieve the nearest images for each query image.\ndef retrieval_results(df_query, df_index, k=21, ef=50, emb_col='embedding'):\n    dim = len(df_index[emb_col].values[0])\n\n    p = hnswlib.Index(space='cosine', dim=dim)\n\n    embeddings = np.vstack(df_index[\"embedding\"].values)\n    p.init_index(max_elements=len(df_index), ef_construction=100, M=16)\n    p.add_items(embeddings)\n    p.set_ef(ef)\n\n    labels, distances = p.knn_query(\n        np.vstack(df_query[\"embedding\"].values),\n        k=k\n    )\n\n    retrieval_rows = []\n\n    for idx, (sub_labels, sub_distances) in enumerate(zip(labels, distances)):\n        gt_image_id = df_query.iloc[idx]['image_id']\n        gt_identity = df_query.iloc[idx]['identity']\n\n        for label_idx, dist in zip(sub_labels, sub_distances):\n            row = df_index.iloc[label_idx]\n\n            retrieved_image_id = row['image_id']\n            retrieved_identity = row['identity']\n\n            # remove self-match safely\n            if retrieved_image_id == gt_image_id:\n                continue\n\n            if (\n                pd.isna(gt_identity)\n                or pd.isna(retrieved_identity)\n            ):\n                is_same_id = None\n            elif gt_identity == retrieved_identity:\n                is_same_id = 1\n            else:\n                is_same_id = 0\n\n            retrieval_rows.append([\n                gt_image_id,\n                gt_identity,\n                retrieved_image_id,\n                retrieved_identity,\n                dist,\n                is_same_id\n            ])\n\n    df_retrieval_results = pd.DataFrame(\n        retrieval_rows,\n        columns=[\n            'main_image_id',\n            'main_identity',\n            'retrieved_image_id',\n            'retrieved_identity',\n            'distance',\n            'is_same_id'\n        ]\n    )\n\n    df_retrieval_results = (\n        df_retrieval_results\n        .sort_values(['main_image_id', 'distance'])\n        .groupby('main_image_id')\n        .head(k - 1)\n        .reset_index(drop=True)\n    )\n\n    return df_retrieval_results","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T22:03:19.832028Z","iopub.execute_input":"2026-05-03T22:03:19.832359Z","iopub.status.idle":"2026-05-03T22:03:19.841771Z","shell.execute_reply.started":"2026-05-03T22:03:19.832333Z","shell.execute_reply":"2026-05-03T22:03:19.840825Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Prepare the images used for retrieval. During test mode, test images are added to the index with unknown identities.\ndef get_data_to_retrieve(df_all, dataset_name, is_test=False):\n    df_sub = df_all[(df_all['dataset'] == dataset_name)].reset_index(drop=True)\n    \n    df_train = df_sub[df_sub['split'] == 'train'].reset_index(drop=True)\n    df_val = df_sub[df_sub['split'] == 'val'].reset_index(drop=True)\n    df_eval = df_sub[df_sub['split'] == 'eval'].reset_index(drop=True)\n\n    if is_test:\n        df_eval['split'] = 'val'\n        df_test = df_sub[df_sub['split'] == 'test'].reset_index(drop=True)\n        df_test['identity'] = None\n\n        df_to_index = pd.concat((df_train, df_val, df_eval, df_test), axis=0, ignore_index=True)\n        df_sub['split'] = df_sub['split'].apply(lambda x: 'val' if x == 'eval' else x)\n    else:\n        df_eval['identity'] = None\n        df_to_index = pd.concat((df_train, df_val, df_eval), axis=0, ignore_index=True)\n\n    return df_sub, df_to_index","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T22:03:22.714706Z","iopub.execute_input":"2026-05-03T22:03:22.715103Z","iopub.status.idle":"2026-05-03T22:03:22.722962Z","shell.execute_reply.started":"2026-05-03T22:03:22.715077Z","shell.execute_reply":"2026-05-03T22:03:22.722021Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load the MiewID model and extract one embedding vector for every image.\nmodel = AutoModel.from_pretrained( \"conservationxlabs/miewid-msv3\", trust_remote_code=True)\nembedding_extractor = DeepFeatures(model=model, device=device, batch_size=128)\nembeddings = embedding_extractor(dataset_all)\n\ndf_all_mixed_split['embedding'] = np.array(embeddings.features).tolist()\n\ndf_all_mixed_split.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Prepare helper functions for batched LightGlue matching. These functions move local features to tensors and match image pairs efficiently.\nimport torch\nimport numpy as np\nfrom tqdm import tqdm\n\n\ndef to_tensor(x, device=None, dtype=None):\n    if torch.is_tensor(x):\n        t = x\n    else:\n        t = torch.tensor(x)\n    if dtype is not None:\n        t = t.to(dtype)\n    if device is not None:\n        t = t.to(device)\n    return t\n\n\ndef prepare_feature_once(feat, device):\n    \"\"\"\n    Converts a single feature dict to tensors once.\n    No batch dimension here.\n    Expected keys:\n      - keypoints: [N, 2]\n      - descriptors: [N, D]\n      - image_size: [2]\n    Optional:\n      - scales: [N]\n      - oris: [N]\n    \"\"\"\n    out = {\n        \"keypoints\": to_tensor(feat[\"keypoints\"], device=device, dtype=torch.float32),\n        \"descriptors\": to_tensor(feat[\"descriptors\"], device=device, dtype=torch.float32),\n        \"image_size\": to_tensor(feat[\"image_size\"], device=device),\n    }\n\n    if \"scales\" in feat:\n        out[\"scales\"] = to_tensor(feat[\"scales\"], device=device, dtype=torch.float32)\n    if \"oris\" in feat:\n        out[\"oris\"] = to_tensor(feat[\"oris\"], device=device, dtype=torch.float32)\n\n    return out\n\n\ndef pad_sequence(list_of_tensors, pad_value=0.0):\n    \"\"\"\n    Pads tensors of shape:\n      [N, D] or [N]\n    into:\n      [B, Nmax, D] or [B, Nmax]\n    \"\"\"\n    max_len = max(x.shape[0] for x in list_of_tensors)\n\n    if list_of_tensors[0].ndim == 2:\n        d = list_of_tensors[0].shape[1]\n        out = torch.full(\n            (len(list_of_tensors), max_len, d),\n            pad_value,\n            dtype=list_of_tensors[0].dtype,\n            device=list_of_tensors[0].device,\n        )\n        for i, x in enumerate(list_of_tensors):\n            out[i, :x.shape[0]] = x\n    elif list_of_tensors[0].ndim == 1:\n        out = torch.full(\n            (len(list_of_tensors), max_len),\n            pad_value,\n            dtype=list_of_tensors[0].dtype,\n            device=list_of_tensors[0].device,\n        )\n        for i, x in enumerate(list_of_tensors):\n            out[i, :x.shape[0]] = x\n    else:\n        raise ValueError(f\"Unsupported ndim={list_of_tensors[0].ndim}\")\n\n    return out\n\n\ndef build_batched_side(feats):\n    \"\"\"\n    feats: list of prepared feature dicts\n    Returns padded batch dict for one side.\n    \"\"\"\n    data = {\n        \"keypoints\": pad_sequence([f[\"keypoints\"] for f in feats], pad_value=0.0),\n        \"descriptors\": pad_sequence([f[\"descriptors\"] for f in feats], pad_value=0.0),\n        \"image_size\": torch.stack([f[\"image_size\"] for f in feats], dim=0),\n    }\n\n    has_scales = all(\"scales\" in f for f in feats)\n    has_oris = all(\"oris\" in f for f in feats)\n\n    if has_scales:\n        data[\"scales\"] = pad_sequence([f[\"scales\"] for f in feats], pad_value=0.0)\n    if has_oris:\n        data[\"oris\"] = pad_sequence([f[\"oris\"] for f in feats], pad_value=0.0)\n\n    return data\n\n\ndef match_feature_batch(matcher, feats0, feats1, ids0, ids1):\n    \"\"\"\n    feats0, feats1: list of prepared features\n    ids0, ids1: list of image ids\n    \"\"\"\n    data0 = build_batched_side(feats0)\n    data1 = build_batched_side(feats1)\n\n    idx0 = torch.tensor(ids0, device=matcher.device)\n    idx1 = torch.tensor(ids1, device=matcher.device)\n\n    batch = (idx0, data0, idx1, data1)\n    return matcher.get_matches(batch)\n\n\ndef batched(iterable, batch_size):\n    for start in range(0, len(iterable), batch_size):\n        yield iterable[start:start + batch_size]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Extract local features for each image and store them by image ID, so they can be reused during matching.\ndef feature_extraction(dataset, feature_extraction_methods):\n    feature_dicts = {}\n    feature_matchers = {}\n    for extractor_name, feature_extractor in feature_extraction_methods:\n        feature_dicts[extractor_name] = {}        \n        \n        features = feature_extractor(dataset)\n        \n        feature_matcher = MatchLightGlue(\n            features=extractor_name,\n            device=device,\n            batch_size=512\n        )\n        \n        feature_matchers[extractor_name] = feature_matcher\n        \n        #for key, feature_dict in feature_dicts.items():\n        for idx, image_id in enumerate(dataset.metadata[\"image_id\"].values):\n            feature_dicts[extractor_name][image_id] = prepare_feature_once(features.features[idx],\n                                                          feature_matcher.device)\n    return feature_dicts, feature_matchers","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Run local feature matching for retrieved image pairs and convert match results into numeric features.\ndef get_match_results(df_retrieval_results, feature_dict, feature_matcher):\n    pairs = []\n    for i, row in df_retrieval_results.iterrows():\n        pairs.append((\n            row[\"main_image_id\"],\n            row[\"retrieved_image_id\"]\n        ))\n    \n    match_results = []\n    \n    batch_size = 256 \n    \n    for batch_pairs in tqdm(list(batched(pairs, batch_size))):\n        ids0 = [p[0] for p in batch_pairs]\n        ids1 = [p[1] for p in batch_pairs]\n    \n        feats0 = [feature_dict[i0] for i0 in ids0]\n        feats1 = [feature_dict[i1] for i1 in ids1]\n    \n        batch_results = match_feature_batch(\n            feature_matcher,\n            feats0,\n            feats1,\n            ids0,\n            ids1\n        )\n        for r in batch_results:\n            match_count = len(r['scores'])\n            lightglue_mean_score = r['scores'].mean()\n            lightglue_match_score = match_count * lightglue_mean_score\n            match_results.append([match_count, lightglue_mean_score, lightglue_match_score])\n            \n    return match_results","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define the clustering function. It uses a trained probability model to decide whether images should join known clusters or create new clusters.\nfrom sklearn.preprocessing import LabelEncoder, StandardScaler\nfrom sklearn.metrics import adjusted_rand_score\nfrom sklearn.pipeline import Pipeline\n\n\ndef run_identity_clustering(df_retrieval_eval, df_sub, model, feature_list, threshold=0.8, top_k=100):\n    dataset_name = df_sub['dataset'].iloc[0] if 'dataset' in df_sub.columns and len(df_sub) > 0 else \"dataset\"\n\n    image_clusters = []\n    cluster_offset = 0\n    generated_identities = {}\n\n    for image_id in df_retrieval_eval['main_image_id'].unique():\n        df_t = df_retrieval_eval[df_retrieval_eval['main_image_id'] == image_id].reset_index(drop=True)\n\n        if len(df_t) == 0:\n            continue\n\n        preds = model.predict_proba(df_t[feature_list])[:, 1]\n\n        pred_list_enrich = list(zip(\n            df_t['retrieved_image_id'].values,\n            df_t['retrieved_identity'].values,\n            preds\n        ))\n        pred_list_enrich = sorted(pred_list_enrich, key=lambda x: x[2], reverse=True)\n\n        main_image_id = df_t['main_image_id'].iloc[0]\n        main_identity = df_t['main_identity'].iloc[0] if 'main_identity' in df_t.columns else None\n\n        if main_image_id in generated_identities:\n            current_main_identity = generated_identities[main_image_id]\n        elif pd.notna(main_identity):\n            current_main_identity = main_identity\n        else:\n            current_main_identity = None\n\n        pred_dict = {}\n        matched_unknown_images = []\n\n        for retrieved_image_id, retrieved_identity, prob in pred_list_enrich[:top_k]:\n            if prob < threshold:\n                continue\n\n            if pd.notna(retrieved_identity):\n                pred_dict[retrieved_identity] = pred_dict.get(retrieved_identity, 0.0) + prob\n            elif retrieved_image_id in generated_identities:\n                assigned_id = generated_identities[retrieved_image_id]\n                pred_dict[assigned_id] = pred_dict.get(assigned_id, 0.0) + prob\n            else:\n                matched_unknown_images.append((retrieved_image_id, prob))\n\n        if current_main_identity is not None:\n            pred_dict[current_main_identity] = pred_dict.get(current_main_identity, 0.0)\n\n        if len(pred_dict) > 0:\n            best_id = max(pred_dict.items(), key=lambda x: x[1])[0]\n        else:\n            best_id = f\"{dataset_name}_new_{cluster_offset}\"\n            cluster_offset += 1\n\n        generated_identities[main_image_id] = best_id\n\n        for retrieved_image_id, _ in matched_unknown_images:\n            generated_identities[retrieved_image_id] = best_id\n\n        image_clusters.append([main_image_id, best_id])\n\n    df_images_clustered = pd.DataFrame(image_clusters, columns=['image_id', 'cluster'])\n    return df_images_clustered, generated_identities\n\n\n\n\n    \ndef evaluate_ari(df_images_clustered, df_sub):\n    df_images_w_gt = pd.merge(\n        df_sub[['image_id', 'identity']],\n        df_images_clustered,\n        on='image_id',\n        how='inner'\n    ).copy()\n\n    df_eval = df_images_w_gt.dropna(subset=['identity', 'cluster']).copy()\n\n    if len(df_eval) == 0:\n        return np.nan, df_eval\n\n    le_gt = LabelEncoder()\n    le_pred = LabelEncoder()\n\n    y_true = le_gt.fit_transform(df_eval['identity'].astype(str))\n    y_pred = le_pred.fit_transform(df_eval['cluster'].astype(str))\n\n    ari = adjusted_rand_score(y_true, y_pred)\n    return ari, df_eval\n\n\ndef tune_threshold_for_ari(df_retrieval_eval, df_sub, model, feature_list, thresholds=None, top_k=100, verbose=False):\n    if thresholds is None:\n        thresholds = np.arange(0.05, 0.96, 0.05)\n\n    results = []\n    best_ari = -np.inf\n    best_threshold = None\n    best_df_clustered = None\n    best_df_eval = None\n\n    for thr in thresholds:\n        df_images_clustered, _ = run_identity_clustering(\n            df_retrieval_eval=df_retrieval_eval,\n            df_sub=df_sub,\n            model=model,\n            feature_list=feature_list,\n            threshold=thr,\n            top_k=top_k\n        )\n\n        ari, df_eval = evaluate_ari(df_images_clustered, df_sub)\n        n_clusters = df_images_clustered['cluster'].nunique() if len(df_images_clustered) > 0 else 0\n\n        results.append({\n            'threshold': thr,\n            'ari': ari,\n            'n_images': len(df_images_clustered),\n            'n_clusters': n_clusters\n        })\n\n        if verbose:\n            print(f\"threshold={thr:.3f} | ARI={ari:.4f} | #clusters={n_clusters}\")\n\n        if pd.notna(ari) and ari > best_ari * 1.01:\n            best_ari = ari\n            best_threshold = thr\n            best_df_clustered = df_images_clustered.copy()\n            best_df_eval = df_eval.copy()\n\n    df_results = pd.DataFrame(results).sort_values('ari', ascending=False).reset_index(drop=True)\n\n    return {\n        'best_threshold': best_threshold,\n        'best_ari': best_ari,\n        'best_clusters': best_df_clustered,\n        'best_eval_df': best_df_eval,\n        'all_results': df_results\n    }\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Feature Matching & Cluster Generation ","metadata":{}},{"cell_type":"code","source":"# Set configuration values: top-k retrieval size, local feature extractors, feature columns, and the binary target column.\nclass cfg:\n    top_k = 21\n    \n    feature_extraction_methods = [['aliked', AlikedExtractor()],\n                       ['sift',SiftExtractor()],\n                       ['superpoint', SuperPointExtractor()],\n                       ['disk', DiskExtractor()]\n                             ]\n    feature_list = [\n        'distance',\n        \n        'aliked_lightglue_match_count',\n        'aliked_lightglue_mean_score',\n        'aliked_lightglue_match_score',\n        \n        'sift_lightglue_match_count',\n        'sift_lightglue_mean_score',\n        'sift_lightglue_match_score',\n    \n        'superpoint_lightglue_match_count',\n        'superpoint_lightglue_mean_score',\n        'superpoint_lightglue_match_score',\n    \n        'disk_lightglue_match_count',\n        'disk_lightglue_mean_score',\n        'disk_lightglue_match_score',\n    ]\n\n    target = 'is_same_id'","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Run the full supervised pipeline for labeled datasets: retrieve neighbors,\n# compute matching features, train logistic regression, tune the threshold, and create test clusters.\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.pipeline import Pipeline\n\ndataset_names = ['SalamanderID2025', 'LynxID2025', 'SeaTurtleID2022']\n\ntest_clusters_dict = {}\nretrieval_dict = {}\nmodel_dict = {}\n\nfor dataset_name in dataset_names:\n\n    df_sub, df_to_index = get_data_to_retrieve(df_all_mixed_split, dataset_name, is_test=True)\n    \n    sub_dataset = WildlifeDataset(\n            df=df_sub, \n            transform=transforms_local\n        )\n    \n    \n    feature_dicts, feature_matchers = feature_extraction(sub_dataset, cfg.feature_extraction_methods)\n    \n    dfs = []\n    for split in df_sub['split'].unique():\n        df = df_sub[df_sub['split'] == split].reset_index(drop=True)\n        df_retrieval = retrieval_results(df, df_to_index)\n        df_retrieval['split'] = split\n        for extractor_name, feature_dict in feature_dicts.items():\n            match_results = get_match_results(df_retrieval, feature_dict, feature_matchers[extractor_name])\n            df_retrieval[[f'{extractor_name}_lightglue_match_count',\n                          f'{extractor_name}_lightglue_mean_score',\n                          f'{extractor_name}_lightglue_match_score']] = match_results\n        dfs.append(df_retrieval)\n    \n    \n    df_retrieval_enriched = pd.concat(dfs, axis=0, ignore_index=True)\n    retrieval_dict[dataset_name] = df_retrieval_enriched\n    \n    df_pairs = df_retrieval_enriched[df_retrieval_enriched['split'].isin(['train', 'val'])].copy()\n    df_pairs = df_pairs[\n        (~df_pairs['retrieved_identity'].isna()) &\n        (~df_pairs['main_identity'].isna())\n    ].reset_index(drop=True)\n    \n    df_pairs = df_pairs.dropna(subset=cfg.feature_list + [cfg.target, 'main_identity', 'main_image_id']).reset_index(drop=True)\n    df_pairs[cfg.target] = df_pairs[cfg.target].astype(int)\n    \n    X = df_pairs[cfg.feature_list].copy()\n    y = df_pairs[cfg.target].copy()\n\n    \n    logreg_final = Pipeline([\n        (\"scaler\", StandardScaler()),\n        (\"model\", LogisticRegression(\n            penalty=\"l2\",\n            C=1.0,\n            class_weight=\"balanced\",\n            random_state=42,\n            solver=\"lbfgs\",\n            max_iter=5000,\n            n_jobs=16,\n        ))\n    ])\n    \n    logreg_final.fit(X, y)\n    model_dict[dataset_name] = logreg_final\n    \n    df_val = df_pairs[df_pairs['split'] == 'val'].copy()\n    df_val['main_identity'] = None\n    df_unknown = df_val[['main_image_id', 'main_identity']].drop_duplicates()\n    unknown_image_ids = dict(zip(df_unknown['main_image_id'], df_unknown['main_identity']))\n    \n    df_val['retrieved_identity'] = df_val.apply(lambda x: unknown_image_ids.get(x['retrieved_image_id'],\n                                                                                x['retrieved_identity']),\n                                                                                 axis=1)\n    \n    threshold_grid = np.arange(0.0, 0.995, 0.05)\n    final_tuning_output = tune_threshold_for_ari(\n        df_retrieval_eval=df_val,\n        df_sub=df_sub,\n        model=logreg_final,\n        feature_list=cfg.feature_list,\n        thresholds=threshold_grid,\n        top_k=cfg.top_k,\n        verbose=True\n    )\n    \n    best_threshold_final = final_tuning_output['best_threshold']\n    best_ari_full = final_tuning_output['best_ari']\n    \n    # metadata for test images\n    df_test_sub = df_sub[df_sub['split'] == 'test'].copy().reset_index(drop=True)\n    print(\"test images:\", df_test_sub.shape)\n    \n    # cluster test images using final trained logistic regression\n    df_test = df_retrieval_enriched[df_retrieval_enriched['split'] == 'test']\n    df_test[cfg.feature_list] = df_test[cfg.feature_list].fillna(df_test[cfg.feature_list].mean())\n    df_test_clusters, generated_identities_test = run_identity_clustering(\n        df_retrieval_eval=df_test,\n        df_sub=df_test_sub,\n        model=logreg_final,\n        feature_list=cfg.feature_list,\n        threshold=best_threshold_final\n    )\n    \n    cluster_assignment = {}\n    c = 0\n    dataset_name = df_test_clusters['cluster'].values[0].split('_')[0]\n    for cluster in df_test_clusters['cluster']:\n        if cluster not in cluster_assignment:\n            cluster_assignment[cluster] = f\"cluster_{dataset_name}_{c}\"\n            c += 1\n    df_test_clusters['mapped_cluster'] = df_test_clusters['cluster'].map(cluster_assignment)\n    df_test_clusters = df_test_clusters.drop('cluster', axis=1)\n    df_test_clusters = df_test_clusters.rename(columns={'mapped_cluster':'cluster'})\n    print(f\"{dataset_name} Cluster Count: {df_test_clusters['cluster'].nunique()}\")\n    test_clusters_dict[dataset_name] = df_test_clusters","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## TexasHornedLizards (Unsupervised Clustering)","metadata":{}},{"cell_type":"code","source":"# Search for a good clustering threshold using silhouette score.\n# This is useful for TexasHornedLizards because it has no labeled training identities.\ndef find_best_threshold_silhouette(X):\n    thresholds = np.arange(0.01, 0.8, 0.02)\n\n    results = []\n\n    for thr in tqdm(thresholds):\n        pred = AgglomerativeClustering(\n            n_clusters=None,\n            metric=\"cosine\",\n            linkage=\"average\",\n            distance_threshold=float(thr)\n        ).fit_predict(X)\n\n        n_clusters = len(np.unique(pred))\n\n        if n_clusters <= 1 or n_clusters >= len(X):\n            continue\n\n        try:\n            score = silhouette_score(X, pred, metric=\"cosine\")\n        except:\n            continue\n\n        results.append({\n            \"threshold\": thr,\n            \"silhouette\": score,\n            \"n_clusters\": n_clusters\n        })\n\n    res_df = pd.DataFrame(results).sort_values(\"silhouette\", ascending=False)\n    return res_df","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Extract TexasHornedLizards embeddings, normalize them, and evaluate silhouette scores over different clustering thresholds.\nfrom sklearn.cluster import AgglomerativeClustering\nfrom sklearn.preprocessing import normalize\nfrom sklearn.metrics import silhouette_score\n\ndataset_name = 'TexasHornedLizards'\n\ndataset_texas = WildlifeDataset(df_all_mixed_split[df_all_mixed_split['dataset'] == dataset_name], transform_miew)\ntexas_embeddings = embedding_extractor(dataset_texas)\n\ndf_texas = df_all_mixed_split[(df_all_mixed_split['dataset'] == dataset_name)].reset_index(drop=True)\nX = np.stack(df_texas[\"embedding\"].values).astype(np.float32)\nX = normalize(X, axis=1)\n\nres_df = find_best_threshold_silhouette(X)\nres_df.head(5)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Plot the number of clusters and silhouette score across thresholds.\n# The dotted line shows the selected threshold.\nimport matplotlib.pyplot as plt\n\nres_df = res_df.sort_values('threshold')\nbest_silhouette_threshold = 0.33\n\nfig, ax1 = plt.subplots()\n\n# Left axis → n_clusters\nax1.plot(res_df['threshold'], res_df['n_clusters'], label='n_clusters')\nax1.set_xlabel('threshold')\nax1.set_ylabel('n_clusters')\n\n# Dotted vertical line at threshold = 0.33\nax1.axvline(\n    x=best_silhouette_threshold,\n    linestyle='--',\n    color='gray',\n    linewidth=1.5,\n    label=f'threshold = {best_silhouette_threshold}'\n)\n\n# Right axis → silhouette\nax2 = ax1.twinx()\nax2.plot(\n    res_df['threshold'],\n    res_df['silhouette'],\n    color='orange',\n    label='silhouette'\n)\nax2.set_ylabel('silhouette')\n\nplt.title('Clusters vs Silhouette over Threshold')\n\n# Combined legend from both axes\nlines1, labels1 = ax1.get_legend_handles_labels()\nlines2, labels2 = ax2.get_legend_handles_labels()\nax1.legend(lines1 + lines2, labels1 + labels2, loc='best')\n\nfig.tight_layout()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Apply agglomerative clustering to TexasHornedLizards using the selected silhouette threshold.\npred = AgglomerativeClustering(\n            n_clusters=None,\n            metric=\"cosine\",\n            linkage=\"average\",\n            distance_threshold=best_silhouette_threshold\n        ).fit_predict(X)\n\ndf_texas[\"cluster_id\"] = pred\ndf_texas[\"cluster\"] = [f\"cluster_{dataset_name}_{i}\" for i in pred]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Combined Submission","metadata":{}},{"cell_type":"code","source":"# Combine clusters from all datasets and save the final submission file.\ndf_cluster_results = pd.concat(list(test_clusters_dict.values()) + [df_texas], axis=0, ignore_index=True)[['image_id', 'cluster']] \nprint(\"Test Size:\",df_cluster_results.shape)\nprint(\"Cluster Count:\", df_cluster_results['cluster'].nunique())\ndf_cluster_results.to_csv('submission.csv', index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}