sem_seg_add.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jupyter environment detected. Enabling Open3D WebVisualizer.\n",
      "[Open3D INFO] WebRTC GUI backend enabled.\n",
      "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import os\n",
    "import random\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import open3d as o3d\n",
    "from sklearn.metrics import accuracy_score, precision_score, recall_score, jaccard_score\n",
    "\n",
    "# Set up device\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "def set_seed(seed=42):\n",
    "    # Set seeds for reproducibility\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "\n",
    "# Set the seed for reproducibility\n",
    "set_seed(42)\n",
    "\n",
    "# Define a linear projection module\n",
    "class LinearProjection(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim):\n",
    "        super(LinearProjection, self).__init__()\n",
    "        self.fc = nn.Linear(input_dim, output_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc(x)\n",
    "\n",
    "# Add the projected CLIP and DINO features\n",
    "def add_features(clip_embeddings, dino_embeddings, clip_projection, dino_projection):\n",
    "    # Project embeddings to a common dimensionality\n",
    "    projected_clip = clip_projection(clip_embeddings)\n",
    "    projected_dino = dino_projection(dino_embeddings)\n",
    "\n",
    "    # Add the projected features\n",
    "    combined_features = projected_clip + projected_dino\n",
    "\n",
    "    return combined_features\n",
    "\n",
    "# Custom Dataset class to handle additive mixture of features\n",
    "class AdditiveScanNetDataset(Dataset):\n",
    "    def __init__(self, scene_ids_file, data_dir, clip_embedding_dir, dino_embedding_dir):\n",
    "        self.scene_ids = self._load_scene_ids(scene_ids_file)\n",
    "        self.data_dir = data_dir\n",
    "        self.clip_embedding_dir = clip_embedding_dir\n",
    "        self.dino_embedding_dir = dino_embedding_dir\n",
    "\n",
    "    def _load_scene_ids(self, scene_ids_file):\n",
    "        with open(scene_ids_file, 'r') as f:\n",
    "            scene_ids = f.read().splitlines()\n",
    "        return scene_ids\n",
    "\n",
    "    def _load_scene_data(self, scene_id):\n",
    "        # Load the .pth file with coordinates, colors, and labels\n",
    "        scene_data = torch.load(os.path.join(self.data_dir, f'{scene_id}_vh_clean_2.pth'))\n",
    "        labels = scene_data[2]\n",
    "\n",
    "        # Load the CLIP and DINO embeddings\n",
    "        clip_data = torch.load(os.path.join(self.clip_embedding_dir, f'{scene_id}.pt'))\n",
    "        clip_embeddings = clip_data['feat']\n",
    "        clip_mask = clip_data['mask_full']\n",
    "\n",
    "        dino_data = torch.load(os.path.join(self.dino_embedding_dir, f'{scene_id}.pt'))\n",
    "        dino_embeddings = dino_data['feat']\n",
    "        dino_mask = dino_data['mask_full']\n",
    "\n",
    "        # Find the intersection of the masks\n",
    "        if not torch.equal(clip_mask, dino_mask):\n",
    "            print(f\"Skipping scene {scene_id} due to mismatched masks.\")\n",
    "            return None\n",
    "        \n",
    "        common_mask = clip_mask & dino_mask\n",
    "        filtered_labels = labels[common_mask]\n",
    "\n",
    "        # Replace label 255 with 20\n",
    "        filtered_labels[filtered_labels == 255] = 20\n",
    "\n",
    "        return clip_embeddings, dino_embeddings, filtered_labels\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.scene_ids)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        scene_id = self.scene_ids[idx]\n",
    "        result = self._load_scene_data(scene_id)\n",
    "\n",
    "        # Skip scenes with mismatched masks\n",
    "        if result is None:\n",
    "            while result is None:\n",
    "                idx = random.randint(0, len(self.scene_ids) - 1)\n",
    "                scene_id = self.scene_ids[idx]\n",
    "                result = self._load_scene_data(scene_id)\n",
    "\n",
    "        clip_embeddings, dino_embeddings, labels = result\n",
    "        return clip_embeddings, dino_embeddings, labels, scene_id\n",
    "\n",
    "# Collate function for variable-sized batches\n",
    "def custom_collate(batch):\n",
    "    clip_list = []\n",
    "    dino_list = []\n",
    "    label_list = []\n",
    "    scene_ids = []\n",
    "\n",
    "    for clip, dino, labels, scene_id in batch:\n",
    "        clip_list.append(clip)\n",
    "        dino_list.append(dino)\n",
    "        label_list.append(labels)\n",
    "        scene_ids.append(scene_id)\n",
    "\n",
    "    return clip_list, dino_list, label_list, scene_ids\n",
    "\n",
    "# Define the segmentation model\n",
    "class SimpleSegmentationModel(nn.Module):\n",
    "    def __init__(self, input_dim, num_classes=21):\n",
    "        super(SimpleSegmentationModel, self).__init__()\n",
    "        self.fc = nn.Linear(input_dim, num_classes)\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc(x.float())\n",
    "        pred = self.softmax(x)\n",
    "        return torch.log(pred)  # Return log-probabilities for NLLLoss\n",
    "\n",
    "# Training function\n",
    "def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=10, patience=3, clip_projection=None, dino_projection=None):\n",
    "    best_val_loss = np.inf\n",
    "    patience_counter = 0\n",
    "    train_losses = []\n",
    "    val_losses = []\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        model.train()\n",
    "        running_loss = 0.0\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        for batch_idx, (clip_list, dino_list, label_list, _) in enumerate(train_loader):\n",
    "            for clip_embeddings, dino_embeddings, labels in zip(clip_list, dino_list, label_list):\n",
    "                clip_embeddings = clip_embeddings.to(device, dtype=torch.float32)\n",
    "                dino_embeddings = dino_embeddings.to(device, dtype=torch.float32)\n",
    "\n",
    "                labels = torch.tensor(labels, dtype=torch.long).to(device)\n",
    "\n",
    "                # Get added features\n",
    "                combined_features = add_features(clip_embeddings, dino_embeddings, clip_projection, dino_projection)\n",
    "\n",
    "                # Forward pass\n",
    "                outputs = model(combined_features)\n",
    "                loss = criterion(outputs, labels)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                running_loss += loss.item()\n",
    "\n",
    "        avg_train_loss = running_loss / len(train_loader)\n",
    "        train_losses.append(avg_train_loss)\n",
    "\n",
    "        # Validation phase\n",
    "        model.eval()\n",
    "        val_loss = 0.0\n",
    "        with torch.no_grad():\n",
    "            for clip_list, dino_list, label_list, _ in val_loader:\n",
    "                for clip_embeddings, dino_embeddings, labels in zip(clip_list, dino_list, label_list):\n",
    "                    clip_embeddings = clip_embeddings.to(device, dtype=torch.float32)\n",
    "                    dino_embeddings = dino_embeddings.to(device, dtype=torch.float32)\n",
    "                    \n",
    "                    labels = torch.tensor(labels, dtype=torch.long).to(device)\n",
    "\n",
    "                    combined_features = add_features(clip_embeddings, dino_embeddings, clip_projection, dino_projection)\n",
    "                    outputs = model(combined_features)\n",
    "                    loss = criterion(outputs, labels)\n",
    "                    val_loss += loss.item()\n",
    "\n",
    "        avg_val_loss = val_loss / len(val_loader)\n",
    "        val_losses.append(avg_val_loss)\n",
    "\n",
    "        print(f\"Epoch [{epoch + 1}/{num_epochs}], Training Loss: {avg_train_loss:.4f}, Validation Loss: {avg_val_loss:.4f}\")\n",
    "\n",
    "        if avg_val_loss < best_val_loss:\n",
    "            best_val_loss = avg_val_loss\n",
    "            patience_counter = 0\n",
    "            torch.save(model.state_dict(), f\"best_model_additive_{epoch}.pth\")\n",
    "        else:\n",
    "            patience_counter += 1\n",
    "\n",
    "        if patience_counter >= patience:\n",
    "            print(f\"Early stopping at epoch {epoch + 1}\")\n",
    "            break\n",
    "\n",
    "# Hyperparameters and settings\n",
    "clip_dim = 3072  # Example dimensions for CLIP\n",
    "dino_dim = 1024  # Example dimensions for DINO\n",
    "target_dim = 1024\n",
    "num_epochs = 10\n",
    "learning_rate = 0.0001\n",
    "\n",
    "# Initialize the projection layers\n",
    "clip_projection = LinearProjection(clip_dim, target_dim).to(device)\n",
    "dino_projection = LinearProjection(dino_dim, target_dim).to(device)\n",
    "\n",
    "# Initialize the model\n",
    "model = SimpleSegmentationModel(input_dim=target_dim, num_classes=21).to(device)  # Updated input_dim to target_dim\n",
    "criterion = nn.NLLLoss()\n",
    "optimizer = optim.Adam(list(model.parameters()) + list(clip_projection.parameters()) + list(dino_projection.parameters()), lr=learning_rate)\n",
    "\n",
    "# Paths to the data\n",
    "train_scene_ids_file = '/projectnb/compvision/charoori/openscene/data/scannet_3d/scannetv2_train_filtered1.txt'\n",
    "val_scene_ids_file = '/projectnb/compvision/charoori/openscene/data/scannet_3d/scannetv2_val_filtered1.txt'\n",
    "train_data_dir = '/projectnb/compvision/charoori/openscene/data/scannet_3d/train'\n",
    "val_data_dir = '/projectnb/compvision/charoori/openscene/data/scannet_3d/val'\n",
    "clip_embedding_dir = '/projectnb/compvision/jteja/Lexicon3D/lexicon3d/clip/clip_features'\n",
    "dino_embedding_dir = '/projectnb/compvision/charoori/Lexicon3D/lexicon3d/dataset/lexicon3d/dinov2_v2/dinov2_features'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skipping scene scene0604_00 due to mismatched masks.\n",
      "Skipping scene scene0646_01 due to mismatched masks.\n",
      "Skipping scene scene0694_00 due to mismatched masks.\n",
      "Skipping scene scene0497_00 due to mismatched masks.\n",
      "Epoch [1/10], Training Loss: 1.1252, Validation Loss: 1.0307\n",
      "Skipping scene scene0646_01 due to mismatched masks.\n",
      "Skipping scene scene0604_00 due to mismatched masks.\n",
      "Skipping scene scene0694_00 due to mismatched masks.\n",
      "Skipping scene scene0497_00 due to mismatched masks.\n",
      "Epoch [2/10], Training Loss: 0.9192, Validation Loss: 0.9663\n",
      "Skipping scene scene0497_00 due to mismatched masks.\n",
      "Skipping scene scene0604_00 due to mismatched masks.\n",
      "Skipping scene scene0646_01 due to mismatched masks.\n",
      "Skipping scene scene0694_00 due to mismatched masks.\n",
      "Epoch [3/10], Training Loss: 0.8608, Validation Loss: 0.9648\n",
      "Skipping scene scene0497_00 due to mismatched masks.\n",
      "Skipping scene scene0646_01 due to mismatched masks.\n",
      "Skipping scene scene0604_00 due to mismatched masks.\n",
      "Skipping scene scene0694_00 due to mismatched masks.\n",
      "Epoch [4/10], Training Loss: 0.8317, Validation Loss: 0.9104\n",
      "Skipping scene scene0646_01 due to mismatched masks.\n",
      "Skipping scene scene0604_00 due to mismatched masks.\n",
      "Skipping scene scene0497_00 due to mismatched masks.\n",
      "Skipping scene scene0694_00 due to mismatched masks.\n",
      "Epoch [5/10], Training Loss: 0.8148, Validation Loss: 0.9377\n",
      "Skipping scene scene0646_01 due to mismatched masks.\n",
      "Skipping scene scene0604_00 due to mismatched masks.\n",
      "Skipping scene scene0497_00 due to mismatched masks.\n",
      "Skipping scene scene0694_00 due to mismatched masks.\n",
      "Epoch [6/10], Training Loss: 0.8024, Validation Loss: 0.9361\n",
      "Skipping scene scene0694_00 due to mismatched masks.\n",
      "Skipping scene scene0694_00 due to mismatched masks.\n",
      "Skipping scene scene0646_01 due to mismatched masks.\n",
      "Skipping scene scene0497_00 due to mismatched masks.\n",
      "Skipping scene scene0604_00 due to mismatched masks.\n",
      "Epoch [7/10], Training Loss: 0.7846, Validation Loss: 0.9282\n",
      "Early stopping at epoch 7\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# Create datasets and data loaders\n",
    "train_dataset = AdditiveScanNetDataset(train_scene_ids_file, train_data_dir, clip_embedding_dir, dino_embedding_dir)\n",
    "val_dataset = AdditiveScanNetDataset(val_scene_ids_file, val_data_dir, clip_embedding_dir, dino_embedding_dir)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True, collate_fn=custom_collate)\n",
    "val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False, collate_fn=custom_collate)\n",
    "\n",
    "# Train the model\n",
    "train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=num_epochs, clip_projection=clip_projection, dino_projection=dino_projection)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.20"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}