{"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,"sourceType":"competition"}],"dockerImageVersionId":31011,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os, random\nfrom collections import defaultdict\nfrom typing import Any\nfrom tqdm import tqdm\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader, RandomSampler\n\nclass ImageDataset(Dataset):\n    def __init__(self, metadata: pd.DataFrame, image_root: str):\n        self.metadata = metadata\n        self.image_root = image_root\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        img_path = os.path.join(self.image_root, row.path)\n        img = Image.open(img_path).convert(\"RGB\")\n        return row.image_id, np.array(img)\n\n\ndef get_dino_processor_and_model(model_name: str) -> tuple:\n    \"\"\"Return processor and model.\"\"\"\n    from transformers import AutoImageProcessor, AutoModel\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n    processor = AutoImageProcessor.from_pretrained(model_name, use_fast=True)\n    model = AutoModel.from_pretrained(model_name).to(device)\n\n    return processor, model\n\n\ndef precompute_embeddings(\n    metadata: pd.DataFrame,\n    image_root: str,\n    processor: Any,\n    model: Any,\n):\n    # we need images in a list, really, don't stack them\n    def custom_collate(batch):\n        ids = [r[0] for r in batch]\n        images = [r[1] for r in batch]\n        return ids, images\n    dataset = ImageDataset(metadata, image_root)\n    dataloader = DataLoader(\n        dataset, \n        batch_size=32, \n        shuffle=False, \n        collate_fn=custom_collate,\n        num_workers=os.cpu_count()\n    )\n\n    res = []\n    with torch.no_grad():\n        for image_ids, images in tqdm(dataloader, mininterval=30):\n            inputs = processor(images=images, return_tensors=\"pt\").to(model.device)\n            outputs = model(**inputs)\n            features = outputs.last_hidden_state\n            embeddings = features[:, 0, :].squeeze(0).cpu().numpy()\n            \n            for image_id, embedding in zip(image_ids, embeddings):\n                res.append({\n                    \"image_id\": image_id,\n                    \"embeddings\": embedding\n                })\n                    \n    return pd.DataFrame(res).sort_values(\"image_id\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"metadata = pd.read_csv('/kaggle/input/animal-clef-2025/metadata.csv')\n\n# pre-compute embeddings\nprocessor, model = get_dino_processor_and_model('facebook/dinov2-base')\nembeddings_df = precompute_embeddings(\n    metadata, '/kaggle/input/animal-clef-2025', processor, model\n)\nembeddings_df.to_parquet(\"embeddings.parquet\")\ndisplay(embeddings_df)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-10T18:46:47.645295Z","iopub.execute_input":"2025-05-10T18:46:47.64586Z","iopub.status.idle":"2025-05-10T18:55:05.518728Z","shell.execute_reply.started":"2025-05-10T18:46:47.645839Z","shell.execute_reply":"2025-05-10T18:55:05.518114Z"}},"outputs":[],"execution_count":null}]}