{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":91451,"databundleVersionId":11223220,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":238987508,"sourceType":"kernelVersion"},{"sourceId":238980352,"sourceType":"kernelVersion"},{"sourceId":383288,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":316386,"modelId":336897}],"dockerImageVersionId":31011,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"! pip install pacmap","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-10T17:21:05.006991Z","iopub.execute_input":"2025-05-10T17:21:05.007780Z","iopub.status.idle":"2025-05-10T17:21:09.054547Z","shell.execute_reply.started":"2025-05-10T17:21:05.007753Z","shell.execute_reply":"2025-05-10T17:21:09.053573Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, random\nfrom collections import defaultdict\nfrom typing import Any\nfrom tqdm import tqdm\nimport pacmap\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader, RandomSampler\n\n\nclass EmbeddingDataset(Dataset):\n    def __init__(self,  metadata: pd.DataFrame):\n        self.metadata = metadata\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n    def __len__(self) -> int:\n        return len(self.metadata)\n\n    def __getitem__(self, idx: int) -> tuple:\n        row = self.metadata.iloc[idx]\n        return row.image_id, torch.from_numpy(row.embeddings).to(self.device)\n\n\nclass ProjectionHead(nn.Module):\n    \"\"\"Simple projection head to transform embeddings.\"\"\"\n\n    def __init__(self, input_dim, output_dim):\n        super(ProjectionHead, self).__init__()\n        self.fc = nn.Linear(input_dim, output_dim)\n        # \n        self.relu = nn.ReLU()\n\n    def forward(self, x):\n        z = self.fc(x)\n        z = self.relu(z)\n        return z","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-10T17:22:39.860063Z","iopub.execute_input":"2025-05-10T17:22:39.860530Z","iopub.status.idle":"2025-05-10T17:22:58.699799Z","shell.execute_reply.started":"2025-05-10T17:22:39.860507Z","shell.execute_reply":"2025-05-10T17:22:58.699247Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dino_embeddings = pd.read_parquet(\"/kaggle/input/preprocess-triplet-embedding/embeddings.parquet\")\nmerged_df = pd.merge(\n    pd.read_csv('/kaggle/input/animal-clef-2025/metadata.csv'),\n    dino_embeddings, \n    on='image_id',\n    how='inner'\n).sort_values(\"image_id\")\ndataset = EmbeddingDataset(merged_df)\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nhead = ProjectionHead(768, 128).to(device)\n\ninput_path = \"/kaggle/input/train-triplet-embedding/head.pt\"\nstate_dict = torch.load(input_path)\nhead.load_state_dict(state_dict)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-10T17:32:47.344669Z","iopub.execute_input":"2025-05-10T17:32:47.345344Z","iopub.status.idle":"2025-05-10T17:32:47.702961Z","shell.execute_reply.started":"2025-05-10T17:32:47.345321Z","shell.execute_reply":"2025-05-10T17:32:47.702095Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dataloader = DataLoader(dataset, batch_size=32, shuffle=False)\n\nres = []\nhead.eval()\nwith torch.no_grad():\n    for batch_image_ids, batch_embeddings in tqdm(dataloader):\n        out = head(batch_embeddings).detach().cpu().numpy()\n        assert out.shape[1] == 128, out.shape\n        for image_id, embedding in zip(batch_image_ids, out):\n            res.append({\n                \"image_id\": int(image_id),\n                \"embeddings\": embedding\n            })\n\ndf = pd.DataFrame(res).sort_values(\"image_id\")\ndisplay(df.head())\ndf.to_parquet(\"embeddings.parquet\", index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-10T17:32:56.524708Z","iopub.execute_input":"2025-05-10T17:32:56.525472Z","iopub.status.idle":"2025-05-10T17:32:58.487282Z","shell.execute_reply.started":"2025-05-10T17:32:56.525447Z","shell.execute_reply":"2025-05-10T17:32:58.486686Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X = np.stack(dino_embeddings.embeddings)\n# 1D embedding to see how the embedding is reshaped\nc = pacmap.PaCMAP(n_components=1).fit_transform(X)\ng = pacmap.PaCMAP().fit_transform(X)\n\nplt.scatter(g[:,0], g[:,1], s=1, alpha=0.5, c=c)\nplt.title(\"DINO embeddings\")\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-10T17:33:05.588561Z","iopub.execute_input":"2025-05-10T17:33:05.589141Z","iopub.status.idle":"2025-05-10T17:33:21.202222Z","shell.execute_reply.started":"2025-05-10T17:33:05.589118Z","shell.execute_reply":"2025-05-10T17:33:21.201479Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# let's generate a plot with pacmap\nX = np.stack(df.embeddings)\ng = pacmap.PaCMAP().fit_transform(X)\n\nplt.scatter(g[:,0], g[:,1], s=1, alpha=0.5, c=c)\nplt.title(\"DINO triplet embeddings\")\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-10T17:33:25.370759Z","iopub.execute_input":"2025-05-10T17:33:25.371321Z","iopub.status.idle":"2025-05-10T17:33:33.793400Z","shell.execute_reply.started":"2025-05-10T17:33:25.371295Z","shell.execute_reply":"2025-05-10T17:33:33.792647Z"}},"outputs":[],"execution_count":null}]}