{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceType":"competition","sourceId":127282,"databundleVersionId":15520105}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import models, transforms\n\n\n# 1\n# User config\n\nMETADATA_CSV = r\"E:\\kg\\animal\\animal-clef-2026\\metadata.csv\"\nIMAGE_ROOT = r\"E:\\kg\\animal\\animal-clef-2026\"\nOUTPUT_CSV = r\"E:\\kg\\animal\\animal-clef-2026\\submission1.csv\"\nMISSING_REPORT_CSV = r\"E:\\kg\\animal\\animal-clef-2026\\missing_paths.csv\"\n\nBATCH_SIZE = 64\nNUM_WORKERS = 0   # safer for notebook / Windows / macOS\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# Set to True only if your pretrained weights can load normally\nUSE_PRETRAINED = False\n\n# Threshold for matching known identities\nKNOWN_SIM_THRESHOLD = 0.55\n\n# Threshold for clustering unknown test samples\nUNKNOWN_CLUSTER_THRESHOLD = 0.70\n\n\n# 2\n# Dataset\nclass AnimalDataset(Dataset):\n    def __init__(self, df, image_root, transform=None):\n        self.df = df.reset_index(drop=True)\n        self.image_root = image_root\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\n        if \"full_path\" in self.df.columns:\n            img_path = row[\"full_path\"]\n        else:\n            img_path = os.path.join(self.image_root, row[\"path\"])\n\n        image = Image.open(img_path).convert(\"RGB\")\n\n        if self.transform is not None:\n            image = self.transform(image)\n\n        identity = \"\"\n        if \"identity\" in self.df.columns and pd.notna(row[\"identity\"]):\n            identity = str(row[\"identity\"]).strip()\n\n        dataset = \"unknown_dataset\"\n        if \"dataset\" in self.df.columns and pd.notna(row[\"dataset\"]):\n            dataset = str(row[\"dataset\"]).strip()\n\n        return {\n            \"image\": image,\n            \"image_id\": int(row[\"image_id\"]),\n            \"identity\": identity,\n            \"dataset\": dataset\n        }\n\n\n# 3\n# Simple union-find for clustering unknown samples\nclass UnionFind:\n    def __init__(self, n):\n        self.parent = list(range(n))\n\n    def find(self, x):\n        while self.parent[x] != x:\n            self.parent[x] = self.parent[self.parent[x]]\n            x = self.parent[x]\n        return x\n\n    def union(self, a, b):\n        ra, rb = self.find(a), self.find(b)\n        if ra != rb:\n            self.parent[rb] = ra\n\n\n# 4\n# Feature extraction\ndef build_model():\n    if USE_PRETRAINED:\n        try:\n            model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)\n            print(\"Loaded pretrained ResNet50 weights.\")\n        except Exception as e:\n            print(f\"Failed to load pretrained weights: {e}\")\n            print(\"Fallback to randomly initialized ResNet50.\")\n            model = models.resnet50(weights=None)\n    else:\n        print(\"Using ResNet50 without pretrained weights.\")\n        model = models.resnet50(weights=None)\n\n    model.fc = torch.nn.Identity()\n    model.eval()\n    return model.to(DEVICE)\n\n\n@torch.no_grad()\ndef extract_features(model, dataloader):\n    all_features = []\n    all_image_ids = []\n    all_identities = []\n    all_datasets = []\n\n    for batch in tqdm(dataloader, desc=\"Extracting features\"):\n        images = batch[\"image\"].to(DEVICE)\n        feats = model(images)\n        feats = F.normalize(feats, p=2, dim=1)\n\n        all_features.append(feats.cpu().numpy())\n\n        if isinstance(batch[\"image_id\"], torch.Tensor):\n            all_image_ids.extend(batch[\"image_id\"].cpu().numpy().tolist())\n        else:\n            all_image_ids.extend(list(batch[\"image_id\"]))\n\n        all_identities.extend(list(batch[\"identity\"]))\n        all_datasets.extend(list(batch[\"dataset\"]))\n\n    all_features = np.concatenate(all_features, axis=0)\n    return all_features, all_image_ids, all_identities, all_datasets\n\n\n# 5\n# Path checking\ndef add_path_check_columns(df):\n    df = df.copy()\n\n    if \"image_id\" not in df.columns:\n        df = df.reset_index().rename(columns={\"index\": \"image_id\"})\n\n    df[\"full_path\"] = df[\"path\"].apply(lambda x: os.path.join(IMAGE_ROOT, x))\n    df[\"path_exists\"] = df[\"full_path\"].apply(os.path.exists)\n\n    missing_df = df.loc[\n        ~df[\"path_exists\"],\n        [\"image_id\", \"path\", \"full_path\", \"split\", \"dataset\"]\n    ].copy()\n\n    print(\"Columns:\", df.columns.tolist())\n    print(\"Total rows:\", len(df))\n    print(\"Number of missing images:\", len(missing_df))\n\n    if len(missing_df) > 0:\n        missing_df.to_csv(MISSING_REPORT_CSV, index=False)\n        print(f\"Missing path report saved to: {MISSING_REPORT_CSV}\")\n        print(\"Examples of missing paths:\")\n        print(missing_df.head(10))\n\n    return df\n\n\n# 6\n# Build prototypes for known train identities\ndef build_identity_prototypes(train_df, train_features):\n    if len(train_df) == 0:\n        return [], np.empty((0, train_features.shape[1]), dtype=np.float32)\n\n    tmp = train_df.copy().reset_index(drop=True)\n    tmp[\"feature_idx\"] = np.arange(len(tmp))\n\n    valid_mask = tmp[\"identity\"].notna() & (tmp[\"identity\"].astype(str).str.strip() != \"\")\n    tmp = tmp[valid_mask].copy()\n\n    if len(tmp) == 0:\n        return [], np.empty((0, train_features.shape[1]), dtype=np.float32)\n\n    identity_to_proto = {}\n    for identity, group in tmp.groupby(\"identity\"):\n        idxs = group[\"feature_idx\"].values\n        proto = train_features[idxs].mean(axis=0)\n        proto = proto / (np.linalg.norm(proto) + 1e-12)\n        identity_to_proto[identity] = proto\n\n    identities = list(identity_to_proto.keys())\n    prototypes = np.stack([identity_to_proto[k] for k in identities], axis=0)\n    return identities, prototypes\n\n\n# 7\n# Assign known / unknown\ndef assign_known_or_unknown(test_features, identities, prototypes, threshold):\n    if len(test_features) == 0:\n        return [], [], np.array([])\n\n    if len(identities) == 0 or len(prototypes) == 0:\n        assigned_known = [None] * len(test_features)\n        unknown_indices = list(range(len(test_features)))\n        return assigned_known, unknown_indices, np.zeros(len(test_features))\n\n    sim = test_features @ prototypes.T\n    best_idx = sim.argmax(axis=1)\n    best_score = sim.max(axis=1)\n\n    assigned_known = []\n    unknown_indices = []\n\n    for i in range(len(test_features)):\n        if best_score[i] >= threshold:\n            assigned_known.append(identities[best_idx[i]])\n        else:\n            assigned_known.append(None)\n            unknown_indices.append(i)\n\n    return assigned_known, unknown_indices, best_score\n\n\n# 8\n# Cluster unknown test samples\ndef cluster_unknowns(test_features_unknown, threshold):\n    n = len(test_features_unknown)\n    if n == 0:\n        return np.array([], dtype=int)\n\n    sim = test_features_unknown @ test_features_unknown.T\n    uf = UnionFind(n)\n\n    for i in range(n):\n        for j in range(i + 1, n):\n            if sim[i, j] >= threshold:\n                uf.union(i, j)\n\n    root_to_cluster = {}\n    cluster_ids = np.zeros(n, dtype=int)\n    next_id = 0\n\n    for i in range(n):\n        root = uf.find(i)\n        if root not in root_to_cluster:\n            root_to_cluster[root] = next_id\n            next_id += 1\n        cluster_ids[i] = root_to_cluster[root]\n\n    return cluster_ids\n\n\n# 9\n# Convert identity to valid submission label\ndef identity_to_submission_label(identity, dataset_name):\n    identity = str(identity).strip()\n    last_part = identity.split(\"_\")[-1]\n\n    if last_part.isdigit():\n        return f\"cluster_{dataset_name}_{last_part}\"\n\n    safe_identity = identity.replace(\" \", \"_\")\n    return f\"cluster_{dataset_name}_{safe_identity}\"\n\n\n# 10\n# Main\ndef main():\n    df = pd.read_csv(METADATA_CSV)\n\n    # Add full_path / path_exists and export missing report\n    df = add_path_check_columns(df)\n\n    train_df = df[df[\"split\"] == \"train\"].copy().reset_index(drop=True)\n    test_df = df[df[\"split\"] == \"test\"].copy().reset_index(drop=True)\n\n    print(f\"Train images (before filtering): {len(train_df)}\")\n    print(f\"Test images (before filtering): {len(test_df)}\")\n    print(f\"Missing in train: {(~train_df['path_exists']).sum()}\")\n    print(f\"Missing in test: {(~test_df['path_exists']).sum()}\")\n\n    # Test images must all exist, otherwise submission row count will be wrong\n    missing_test_count = (~test_df[\"path_exists\"]).sum()\n    if missing_test_count > 0:\n        raise FileNotFoundError(\n            f\"There are {missing_test_count} missing test images. \"\n            f\"Please fix paths first. See: {MISSING_REPORT_CSV}\"\n        )\n\n    # Drop missing train images only\n    train_df = train_df[train_df[\"path_exists\"]].copy().reset_index(drop=True)\n    test_df = test_df[test_df[\"path_exists\"]].copy().reset_index(drop=True)\n\n    print(f\"Train images (after filtering): {len(train_df)}\")\n    print(f\"Test images (after filtering): {len(test_df)}\")\n\n    transform = transforms.Compose([\n        transforms.Resize((224, 224)),\n        transforms.ToTensor(),\n        transforms.Normalize(\n            mean=[0.485, 0.456, 0.406],\n            std=[0.229, 0.224, 0.225]\n        )\n    ])\n\n    train_dataset = AnimalDataset(train_df, IMAGE_ROOT, transform=transform)\n    test_dataset = AnimalDataset(test_df, IMAGE_ROOT, transform=transform)\n\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=BATCH_SIZE,\n        shuffle=False,\n        num_workers=NUM_WORKERS,\n        pin_memory=(DEVICE == \"cuda\")\n    )\n\n    test_loader = DataLoader(\n        test_dataset,\n        batch_size=BATCH_SIZE,\n        shuffle=False,\n        num_workers=NUM_WORKERS,\n        pin_memory=(DEVICE == \"cuda\")\n    )\n\n    model = build_model()\n\n    train_features, _, _, _ = extract_features(model, train_loader)\n    test_features, _, _, _ = extract_features(model, test_loader)\n\n    # Final prediction container\n    pred_clusters = np.empty(len(test_df), dtype=object)\n\n    # Process each dataset separately\n    dataset_list = sorted(test_df[\"dataset\"].dropna().astype(str).unique().tolist())\n    print(\"Datasets in test:\", dataset_list)\n\n    for dataset_name in dataset_list:\n        print(f\"\\nProcessing dataset: {dataset_name}\")\n\n        train_mask = (train_df[\"dataset\"].astype(str) == dataset_name).to_numpy()\n        test_mask = (test_df[\"dataset\"].astype(str) == dataset_name).to_numpy()\n\n        train_df_d = train_df.loc[train_mask].copy().reset_index(drop=True)\n        test_df_d = test_df.loc[test_mask].copy().reset_index(drop=True)\n\n        train_features_d = train_features[train_mask]\n        test_features_d = test_features[test_mask]\n\n        print(f\"  Train images: {len(train_df_d)}\")\n        print(f\"  Test images: {len(test_df_d)}\")\n\n        identities, prototypes = build_identity_prototypes(train_df_d, train_features_d)\n        print(f\"  Known train identities: {len(identities)}\")\n\n        assigned_known, unknown_indices, _ = assign_known_or_unknown(\n            test_features_d,\n            identities,\n            prototypes,\n            threshold=KNOWN_SIM_THRESHOLD\n        )\n\n        print(f\"  Matched to known identities: {sum(x is not None for x in assigned_known)}\")\n        print(f\"  Unknown test samples: {len(unknown_indices)}\")\n\n        unknown_features = test_features_d[unknown_indices]\n        unknown_cluster_ids = cluster_unknowns(\n            unknown_features,\n            threshold=UNKNOWN_CLUSTER_THRESHOLD\n        )\n\n        # Start numbering new clusters after existing numeric suffixes\n        existing_numeric_suffix = []\n        for identity in identities:\n            try:\n                existing_numeric_suffix.append(int(str(identity).split(\"_\")[-1]))\n            except Exception:\n                pass\n\n        if len(existing_numeric_suffix) > 0:\n            next_new_cluster_id = max(existing_numeric_suffix) + 1\n        else:\n            next_new_cluster_id = 0\n\n        unknown_local_to_global = {}\n        pred_clusters_d = []\n        unknown_ptr = 0\n\n        for i in range(len(test_df_d)):\n            if assigned_known[i] is not None:\n                pred_clusters_d.append(\n                    identity_to_submission_label(assigned_known[i], dataset_name)\n                )\n            else:\n                local_cluster = int(unknown_cluster_ids[unknown_ptr])\n\n                if local_cluster not in unknown_local_to_global:\n                    unknown_local_to_global[local_cluster] = next_new_cluster_id\n                    next_new_cluster_id += 1\n\n                global_cluster = unknown_local_to_global[local_cluster]\n                pred_clusters_d.append(f\"cluster_{dataset_name}_{global_cluster}\")\n                unknown_ptr += 1\n\n        pred_clusters[test_mask] = pred_clusters_d\n\n    # Build final submission\n    submission = pd.DataFrame({\n        \"image_id\": test_df[\"image_id\"].astype(int),\n        \"cluster\": pred_clusters\n    })\n\n    # Validate prefixes\n    valid_prefixes = (\n        \"cluster_LynxID2025_\",\n        \"cluster_SalamanderID2025_\",\n        \"cluster_SeaTurtleID2022_\",\n        \"cluster_TexasHornedLizards_\"\n    )\n\n    bad_mask = ~submission[\"cluster\"].astype(str).str.startswith(valid_prefixes)\n    bad_count = bad_mask.sum()\n\n    print(\"\\nIllegal label count:\", bad_count)\n    if bad_count > 0:\n        print(submission.loc[bad_mask].head(20))\n        raise ValueError(\"Some labels do not have valid competition prefixes.\")\n\n    submission.to_csv(OUTPUT_CSV, index=False)\n    print(f\"\\nSaved submission to: {OUTPUT_CSV}\")\n    print(submission.head(20))\n\n\nif __name__ == \"__main__\":\n    main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}