{"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":31012,"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","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:14.271658Z","iopub.execute_input":"2025-05-08T04:47:14.272310Z","iopub.status.idle":"2025-05-08T04:47:14.276672Z","shell.execute_reply.started":"2025-05-08T04:47:14.272284Z","shell.execute_reply":"2025-05-08T04:47:14.275880Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"all_metadata = pd.read_csv('/kaggle/input/animal-clef-2025/metadata.csv')\nmetadata = all_metadata.loc[all_metadata['split'] == 'database']\nmetadata","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:14.925312Z","iopub.execute_input":"2025-05-08T04:47:14.925880Z","iopub.status.idle":"2025-05-08T04:47:14.971415Z","shell.execute_reply.started":"2025-05-08T04:47:14.925857Z","shell.execute_reply":"2025-05-08T04:47:14.970720Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dfs = []\nfor identity in metadata.identity.unique():\n    id_metadata = metadata.loc[metadata.identity == identity].copy()\n    \n    if id_metadata.shape[0] >= 2:\n        dfs.append(id_metadata)\n\nmetadata = pd.concat(dfs, axis=0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:15.444337Z","iopub.execute_input":"2025-05-08T04:47:15.444982Z","iopub.status.idle":"2025-05-08T04:47:16.709477Z","shell.execute_reply.started":"2025-05-08T04:47:15.444962Z","shell.execute_reply":"2025-05-08T04:47:16.708957Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"metadata","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:18.076355Z","iopub.execute_input":"2025-05-08T04:47:18.076966Z","iopub.status.idle":"2025-05-08T04:47:18.087164Z","shell.execute_reply.started":"2025-05-08T04:47:18.076942Z","shell.execute_reply":"2025-05-08T04:47:18.086480Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:20.474723Z","iopub.execute_input":"2025-05-08T04:47:20.475294Z","iopub.status.idle":"2025-05-08T04:47:20.479768Z","shell.execute_reply.started":"2025-05-08T04:47:20.475270Z","shell.execute_reply":"2025-05-08T04:47:20.478866Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class EmbeddingDataset(Dataset):\n    def __init__(\n        self,\n        metadata: pd.DataFrame,\n        img_dir: str,\n        processor: Any,\n        model: Any,\n        embed_type: str = 'cls',\n        target_col: str = 'identity',\n        img_id_col: str = 'image_id'\n    ):\n\n        self.metadata = metadata\n        self.img_dir = img_dir\n        self.processor = processor\n        self.model = model\n        self.embed_type = embed_type\n        self.target_col = target_col\n        self.img_id_col = img_id_col\n        self.device = next(model.parameters()).device\n\n        self.classes = sorted(metadata[target_col].astype(str).unique().tolist())\n        self.class_to_idx = {c: i for i, c in enumerate(self.classes)}\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\n        img_id = row[self.img_id_col]\n        img_path = os.path.join(self.img_dir, row['path'])\n        img = Image.open(img_path).convert('RGB')\n\n        with torch.no_grad():\n            inputs = self.processor(images=img, return_tensors='pt').to(self.device)\n            outputs = self.model(**inputs)\n            features = outputs.last_hidden_state\n\n            if self.embed_type == 'cls':\n                embedding = features[:, 0, :].squeeze(0)\n            else:\n                embedding = features[:, 1:, :].mean(dim=1).squeeze(0)\n\n        target = row[self.target_col]\n        target_int = self.class_to_idx[target]\n        label = torch.tensor(target_int, dtype=torch.long)\n\n        return embedding, label, target","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:21.137171Z","iopub.execute_input":"2025-05-08T04:47:21.137433Z","iopub.status.idle":"2025-05-08T04:47:21.144924Z","shell.execute_reply.started":"2025-05-08T04:47:21.137412Z","shell.execute_reply":"2025-05-08T04:47:21.144208Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TripletLoss(nn.Module):\n    \"\"\"\n    Triplet loss with a margin.\n    Takes embeddings of an anchor, a positive and a negative sample.\n    \"\"\"\n    \n    def __init__(self, margin=1.0):\n        super(TripletLoss, self).__init__()\n        self.margin = margin\n        \n    def forward(self, anchor, positive, negative):\n        # Calculate distances\n        pos_dist = torch.pairwise_distance(anchor, positive)\n        neg_dist = torch.pairwise_distance(anchor, negative)\n        \n        # Calculate triplet loss\n        loss = torch.relu(pos_dist - neg_dist + self.margin)\n        \n        return loss.mean()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:24.118568Z","iopub.execute_input":"2025-05-08T04:47:24.119280Z","iopub.status.idle":"2025-05-08T04:47:24.124078Z","shell.execute_reply.started":"2025-05-08T04:47:24.119255Z","shell.execute_reply":"2025-05-08T04:47:24.123278Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ProjectionHead(nn.Module):\n    \"\"\"\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    def forward(self, x):\n        return self.fc(x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:26.877946Z","iopub.execute_input":"2025-05-08T04:47:26.878230Z","iopub.status.idle":"2025-05-08T04:47:26.882626Z","shell.execute_reply.started":"2025-05-08T04:47:26.878209Z","shell.execute_reply":"2025-05-08T04:47:26.882071Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_with_online_triplet_generation(\n    dataset,\n    model=None,  # Not needed as embeddings are precomputed in this dataset\n    batch_size=32,\n    epochs=10,\n    margin=1.0,\n    learning_rate=0.001,\n    triplets_per_anchor=1,\n    projection_head=None\n):\n    \"\"\"\n    Train with online triplet generation using a standard DataLoader.\n    \n    Args:\n        dataset: Instance of EmbeddingDataset\n        model: The model to train (if needed)\n        batch_size: Batch size for DataLoader\n        epochs: Number of training epochs\n        margin: Margin for triplet loss\n        learning_rate: Learning rate for optimizer\n        triplets_per_anchor: Number of triplets to generate per anchor\n        projection_head: Optional projection head to train on top of embeddings\n    \"\"\"\n    # Create standard DataLoader with RandomSampler\n    dataloader = DataLoader(\n        dataset,\n        batch_size=batch_size,\n        sampler=RandomSampler(dataset),\n    )\n    \n    # Create a label to indices mapping\n    label_to_indices = defaultdict(list)\n    for idx in tqdm(range(len(dataset))):\n        _, label, _ = dataset[idx]\n        label_to_indices[label.item()].append(idx)\n    \n    # Initialize projection head if provided\n    if projection_head is not None:\n        optimizer = optim.Adam(projection_head.parameters(), lr=learning_rate)\n        device = next(projection_head.parameters()).device\n    else:\n        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \n    # Initialize triplet loss\n    criterion = TripletLoss(margin=margin)\n    \n    # Training loop\n    for epoch in range(epochs):\n        running_loss = 0.0\n        triplet_count = 0\n        \n        # Process batches\n        for batch_embeddings, batch_labels, _ in tqdm(dataloader):\n            batch_size = batch_embeddings.size(0)\n            \n            # Skip small batches\n            if batch_size <= 1:\n                continue\n            \n            if projection_head is not None:\n                # Forward through projection head if provided\n                batch_embeddings = projection_head(batch_embeddings)\n            \n            # Process each anchor in the batch\n            batch_loss = 0\n            valid_triplets = 0\n            \n            for i in range(batch_size):\n                anchor_embedding = batch_embeddings[i].unsqueeze(0)\n                anchor_label = batch_labels[i].item()\n                \n                # Find all positive indices (same class as anchor)\n                # Excluding the anchor itself\n                positive_indices = [idx for idx in label_to_indices[anchor_label] if idx != i]\n                \n                # Find all negative indices (different class from anchor)\n                negative_classes = [label for label in label_to_indices.keys() if label != anchor_label]\n                \n                # Skip if no valid triplets can be formed\n                if not positive_indices or not negative_classes:\n                    continue\n                \n                # Generate multiple triplets per anchor\n                for _ in range(min(triplets_per_anchor, len(positive_indices))):\n                    # Randomly select a positive\n                    pos_idx = random.choice(positive_indices)\n                    pos_embedding, _, _ = dataset[pos_idx]\n                    pos_embedding = pos_embedding.unsqueeze(0)\n                    \n                    # Randomly select a negative class and then a negative example\n                    neg_class = random.choice(negative_classes)\n                    neg_idx = random.choice(label_to_indices[neg_class])\n                    neg_embedding, _, _ = dataset[neg_idx]\n                    neg_embedding = neg_embedding.unsqueeze(0)\n                    \n                    # Apply projection head if needed\n                    if projection_head is not None:\n                        pos_embedding = projection_head(pos_embedding)\n                        neg_embedding = projection_head(neg_embedding)\n                    \n                    # Compute loss for this triplet\n                    triplet_loss = criterion(anchor_embedding, pos_embedding, neg_embedding)\n                    batch_loss += triplet_loss\n                    valid_triplets += 1\n            \n            # Skip batch if no valid triplets\n            if valid_triplets == 0:\n                continue\n            \n            # Normalize loss by number of triplets\n            batch_loss = batch_loss / valid_triplets\n            \n            # Backward and optimize if we have a projection head\n            if projection_head is not None:\n                optimizer.zero_grad()\n                batch_loss.backward()\n                optimizer.step()\n            \n            # Update statistics\n            running_loss += batch_loss.item()\n            triplet_count += 1\n        \n        # Print epoch statistics\n        if triplet_count > 0:\n            epoch_loss = running_loss / triplet_count\n            print(f'Epoch {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}')\n        else:\n            print(f'Epoch {epoch+1}/{epochs}, No valid triplets found')\n    \n    return projection_head if projection_head is not None else None","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:27.438274Z","iopub.execute_input":"2025-05-08T04:47:27.438773Z","iopub.status.idle":"2025-05-08T04:47:27.450001Z","shell.execute_reply.started":"2025-05-08T04:47:27.438752Z","shell.execute_reply":"2025-05-08T04:47:27.449259Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"processor, model = get_dino_processor_and_model('facebook/dinov2-base')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T00:32:24.022570Z","iopub.execute_input":"2025-05-08T00:32:24.022735Z","iopub.status.idle":"2025-05-08T00:32:24.569217Z","shell.execute_reply.started":"2025-05-08T00:32:24.022722Z","shell.execute_reply":"2025-05-08T00:32:24.568654Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dataset = EmbeddingDataset(metadata, '/kaggle/input/animal-clef-2025', processor, model)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:34.896415Z","iopub.execute_input":"2025-05-08T04:47:34.897113Z","iopub.status.idle":"2025-05-08T04:47:34.901755Z","shell.execute_reply.started":"2025-05-08T04:47:34.897092Z","shell.execute_reply":"2025-05-08T04:47:34.901107Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"proj_head = ProjectionHead(768, 128).to(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nproj_head","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T00:32:24.575638Z","iopub.execute_input":"2025-05-08T00:32:24.575912Z","iopub.status.idle":"2025-05-08T00:32:24.591624Z","shell.execute_reply.started":"2025-05-08T00:32:24.575887Z","shell.execute_reply":"2025-05-08T00:32:24.590994Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"head = train_with_online_triplet_generation(dataset, projection_head=proj_head)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T00:32:24.592267Z","iopub.execute_input":"2025-05-08T00:32:24.592426Z","iopub.status.idle":"2025-05-08T03:44:18.174854Z","shell.execute_reply.started":"2025-05-08T00:32:24.592414Z","shell.execute_reply":"2025-05-08T03:44:18.174067Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.save(head.state_dict(), '/kaggle/working/head.pt')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T03:44:18.175618Z","iopub.execute_input":"2025-05-08T03:44:18.175826Z","iopub.status.idle":"2025-05-08T03:44:18.184545Z","shell.execute_reply.started":"2025-05-08T03:44:18.175789Z","shell.execute_reply":"2025-05-08T03:44:18.184072Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"state_dict = torch.load('/kaggle/working/head.pt')\nhead = ProjectionHead(768, 128).to(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nhead.load_state_dict(state_dict)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:39.478425Z","iopub.execute_input":"2025-05-08T04:47:39.478971Z","iopub.status.idle":"2025-05-08T04:47:39.488286Z","shell.execute_reply.started":"2025-05-08T04:47:39.478951Z","shell.execute_reply":"2025-05-08T04:47:39.487575Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dataloader = DataLoader(\n    dataset,\n    batch_size=32,\n    shuffle=False\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T04:47:59.081339Z","iopub.execute_input":"2025-05-08T04:47:59.082106Z","iopub.status.idle":"2025-05-08T04:47:59.085675Z","shell.execute_reply.started":"2025-05-08T04:47:59.082074Z","shell.execute_reply":"2025-05-08T04:47:59.084976Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"embeddings_df = []","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T05:08:43.306497Z","iopub.execute_input":"2025-05-08T05:08:43.306787Z","iopub.status.idle":"2025-05-08T05:08:43.362276Z","shell.execute_reply.started":"2025-05-08T05:08:43.306759Z","shell.execute_reply":"2025-05-08T05:08:43.361539Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"head.eval()\nwith torch.no_grad():\n    for batch_embeddings, batch_labels, batch_targets in tqdm(dataloader):\n        out_embeddings = head(batch_embeddings).detach().cpu().numpy()\n        out_labels = batch_labels.detach().cpu().numpy()\n        out_targets = np.array(batch_targets)\n\n        out = np.concatenate((out_targets.reshape(-1, 1), out_labels.reshape(-1, 1), out_embeddings), axis=1)\n        if out.shape != (32, 130):\n            print(f'Error! out.shape = {out.shape}')\n            \n        embeddings_df.append(out)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T05:08:45.908532Z","iopub.execute_input":"2025-05-08T05:08:45.909232Z","iopub.status.idle":"2025-05-08T05:14:33.741399Z","shell.execute_reply.started":"2025-05-08T05:08:45.909207Z","shell.execute_reply":"2025-05-08T05:14:33.740580Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"embeddings_df = np.concatenate(embeddings_df, axis=0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T05:16:12.771359Z","iopub.execute_input":"2025-05-08T05:16:12.771900Z","iopub.status.idle":"2025-05-08T05:16:12.947716Z","shell.execute_reply.started":"2025-05-08T05:16:12.771875Z","shell.execute_reply":"2025-05-08T05:16:12.946733Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"columns = ['identity', 'label'] + [f'embed_dim{i}' for i in range(len(embeddings_df[0]) - 2)]\nembeddings_df = pd.DataFrame(embeddings_df, columns=columns)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T05:16:34.497824Z","iopub.execute_input":"2025-05-08T05:16:34.498476Z","iopub.status.idle":"2025-05-08T05:16:34.672499Z","shell.execute_reply.started":"2025-05-08T05:16:34.498454Z","shell.execute_reply":"2025-05-08T05:16:34.671947Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"embeddings_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T05:16:39.371106Z","iopub.execute_input":"2025-05-08T05:16:39.371549Z","iopub.status.idle":"2025-05-08T05:16:39.421070Z","shell.execute_reply.started":"2025-05-08T05:16:39.371529Z","shell.execute_reply":"2025-05-08T05:16:39.420337Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"embeddings_df.to_csv('/kaggle/working/triplet_dataset_embeddings.csv', index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T05:17:12.452348Z","iopub.execute_input":"2025-05-08T05:17:12.452603Z","iopub.status.idle":"2025-05-08T05:17:13.226354Z","shell.execute_reply.started":"2025-05-08T05:17:12.452585Z","shell.execute_reply":"2025-05-08T05:17:13.225748Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}