DINO_localization_vindr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "826ba406-4d30-4588-819f-423b5eded5b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0306e9c6-01ec-4993-9e88-8dd11b447a85",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in /home/atricham/.cache/torch/hub/facebookresearch_dino_main\n"
     ]
    }
   ],
   "source": [
    "vitb16 = torch.hub.load('facebookresearch/dino:main', 'dino_vitb16')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "50df29e3-ae88-4624-bcc6-f2eaab9eaba4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = vitb16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "afcfda38-0428-4493-845a-e9c02c1e9d5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in /home/atricham/.cache/torch/hub/facebookresearch_dino_main\n",
      "/home/atricham/.local/lib/python3.9/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "/home/atricham/.local/lib/python3.9/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=None`.\n",
      "  warnings.warn(msg)\n"
     ]
    }
   ],
   "source": [
    "resnet50 = torch.hub.load('facebookresearch/dino:main', 'dino_resnet50')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f0274c7e-5a54-4903-b4fb-95bb205beb58",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = resnet50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c75d3c8b-c29d-443c-8094-04affac32f7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "\n",
    "# Define a custom dataset class\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, csv_file, image_dir, transform=None):\n",
    "        self.data = pd.read_csv(csv_file)\n",
    "        self.image_dir = image_dir\n",
    "        self.transform = transform\n",
    "\n",
    "        # Map unique class labels to integer values\n",
    "        unique_classes = self.data['class_name'].unique()\n",
    "        self.class_to_int = {class_name: idx for idx, class_name in enumerate(unique_classes)}\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        image_id = self.data.iloc[idx, 0]  # 'image_id' is the first column\n",
    "        image_filename = f\"{image_id}.jpg\"  # Add \".jpg\" extension\n",
    "        image_path = os.path.join(self.image_dir, image_filename)\n",
    "        image = Image.open(image_path).convert(\"RGB\")\n",
    "\n",
    "        # Extract width and height\n",
    "        class_name = self.data.iloc[idx, 1]\n",
    "        width = self.data.iloc[idx, 8]  # 'width' is in the ninth column\n",
    "        height = self.data.iloc[idx, 9]  # 'height' is in the tenth column\n",
    "\n",
    "        # Extract bounding box coordinates or calculate them if missing\n",
    "        x_min = self.data.iloc[idx, 4]\n",
    "        y_min = self.data.iloc[idx, 5]\n",
    "        x_max = self.data.iloc[idx, 6]\n",
    "        y_max = self.data.iloc[idx, 7]\n",
    "\n",
    "        if pd.isna(x_min) or pd.isna(y_min) or pd.isna(x_max) or pd.isna(y_max):\n",
    "            # Calculate bounding box coordinates as a percentage of image dimensions\n",
    "            x_min = 0\n",
    "            y_min = 0\n",
    "            x_max = 1.0  # Set x_max to 100% of image width\n",
    "            y_max = 1.0  # Set y_max to 100% of image height\n",
    "\n",
    "        # bbox = [x_min, y_min, x_max, y_max]\n",
    "        bbox = torch.tensor([x_min, y_min, x_max, y_max], dtype=torch.float32)\n",
    "\n",
    "\n",
    "        # Map class name to integer label\n",
    "        class_label = self.class_to_int[class_name]\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "            \n",
    "        # Return data as a dictionary\n",
    "        data_dict = {\n",
    "            'images': image,\n",
    "            'targets': {\n",
    "                'class_label': class_label,\n",
    "                'bbox': bbox,\n",
    "                'width': width,\n",
    "                'height': height\n",
    "            }\n",
    "        }\n",
    "\n",
    "        return data_dict\n",
    "\n",
    "# Define a custom collate function for the DataLoader\n",
    "def custom_collate_fn(batch):\n",
    "    images = [item['images'] for item in batch]\n",
    "    targets = [item['targets'] for item in batch]\n",
    "\n",
    "    # Convert images and targets into a batch\n",
    "    images = torch.stack(images, dim=0)\n",
    "    class_labels = torch.tensor([item['class_label'] for item in targets], dtype=torch.long)\n",
    "    bboxes = torch.stack([item['bbox'] for item in targets], dim=0) \n",
    "    widths = torch.tensor([item['width'] for item in targets], dtype=torch.float32)\n",
    "    heights = torch.tensor([item['height'] for item in targets], dtype=torch.float32)\n",
    "\n",
    "    return {\n",
    "        'images': images,\n",
    "        'class_labels': class_labels,\n",
    "        'bboxes': bboxes,\n",
    "        'widths': widths,\n",
    "        'heights': heights\n",
    "    }\n",
    "\n",
    "# Specify the path to your CSV file and image directory\n",
    "csv_file = 'train.csv'  # Replace with the actual path to your CSV file\n",
    "image_dir = 'dataset-jpg/train'  # Replace with the actual path to your image directory\n",
    "\n",
    "# Define any image transformations you want to apply (e.g., resizing, normalization)\n",
    "transform = transforms.Compose([transforms.Resize((224, 224)),\n",
    "                                transforms.ToTensor()])\n",
    "\n",
    "# Create an instance of the custom dataset\n",
    "custom_dataset = CustomDataset(csv_file, image_dir, transform=transform)\n",
    "\n",
    "# Create a data loader with custom collate function\n",
    "batch_size = 8  # Adjust the batch size according to your needs\n",
    "train_data_loader = DataLoader(custom_dataset, batch_size=batch_size, shuffle=True, collate_fn=custom_collate_fn)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7453fc7b-b2fa-49d1-b6a6-85466a05c237",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "\n",
    "# Define a custom dataset class for the test data\n",
    "class TestDataset(Dataset):\n",
    "    def __init__(self, csv_file, image_dir, transform=None):\n",
    "        self.data = pd.read_csv(csv_file)\n",
    "        self.image_dir = image_dir\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        image_id = self.data.iloc[idx, 0]  # 'image_id' is the first column\n",
    "        image_filename = f\"{image_id}.jpg\"  # Add \".jpg\" extension\n",
    "        image_path = os.path.join(self.image_dir, image_filename)\n",
    "        image = Image.open(image_path).convert(\"RGB\")\n",
    "\n",
    "        # Extract width and height\n",
    "        width = self.data.iloc[idx, 1]  # 'width' is in the second column\n",
    "        height = self.data.iloc[idx, 2]  # 'height' is in the third column\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "\n",
    "        # Create targets dictionary\n",
    "        targets = {\n",
    "            'width': width,\n",
    "            'height': height\n",
    "        }\n",
    "\n",
    "        return image, targets\n",
    "\n",
    "# Specify the path to your test CSV file and image directory\n",
    "test_csv_file = 'test.csv'  # Replace with the actual path to your test CSV file\n",
    "test_image_dir = 'dataset-jpg/test'  # Replace with the actual path to your test image directory\n",
    "\n",
    "# Define any image transformations you want to apply (e.g., resizing, normalization)\n",
    "transform = transforms.Compose([transforms.Resize((224, 224)),\n",
    "                                transforms.ToTensor()])\n",
    "\n",
    "# Create an instance of the custom test dataset\n",
    "test_dataset = TestDataset(test_csv_file, test_image_dir, transform=transform)\n",
    "\n",
    "# Create a data loader for the test data\n",
    "batch_size = 8  # Adjust the batch size according to your needs\n",
    "test_data_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "# Now you can load data using the test_data_loader and pass it to your model for bounding box detection\n",
    "# To mount the data to a device, you can do something like this:\n",
    "# for images, targets in test_data_loader:\n",
    "#     images = images.to(device)\n",
    "#     targets = {key: value.to(device) for key, value in targets.items()}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dc2e4652-63a9-4058-a98a-efa9cc6a5352",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResNet(\n",
       "  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (layer1): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer2): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (3): Bottleneck(\n",
       "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer3): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (3): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (4): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (5): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer4): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (fc): Identity()\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_epochs = 3\n",
    "# Define the device (CPU or GPU)\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4e50bfa2-828d-4992-ad73-a5b3d14d258a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Define the number of object classes and bounding box parameters\n",
    "num_classes = 14  # Replace with the number of object classes in your dataset\n",
    "num_bbox_params = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "a5d45683-c7f1-4207-b711-f92655165cac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of output channels for class_logits: 14\n",
      "Number of output channels for bbox_pred: 4\n"
     ]
    }
   ],
   "source": [
    "# Define a custom object detection model\n",
    "class ObjectDetectionModel(nn.Module):\n",
    "    def __init__(self, backbone_model, num_classes, num_bbox_params=4):\n",
    "        super(ObjectDetectionModel, self).__init__()\n",
    "\n",
    "        # Load the ResNet-50 backbone\n",
    "        self.resnet_backbone = backbone_model\n",
    "\n",
    "        # Object classification head\n",
    "        self.cls_head = nn.Sequential(\n",
    "            nn.Linear(2048, 512),  # Adjust input features if needed\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, num_classes)  # Output logits for each class\n",
    "        )\n",
    "\n",
    "        # Bounding box regression head\n",
    "        self.bbox_head = nn.Sequential(\n",
    "            nn.Linear(2048, 512),  # Adjust input features if needed\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, num_bbox_params)  # Output bounding box parameters\n",
    "        )\n",
    "\n",
    "        def forward(self, x):\n",
    "            # Backbone feature extraction\n",
    "            x = self.resnet_backbone(x)\n",
    "\n",
    "            # Object classification\n",
    "            class_logits = self.cls_head(x)\n",
    "\n",
    "            # Bounding box regression\n",
    "            bbox_pred = self.bbox_head(x)\n",
    "\n",
    "            return class_logits, bbox_pred\n",
    "\n",
    "\n",
    "\n",
    "# Instantiate your custom object detection model with the ResNet-50 backbone\n",
    "object_detection_model = ObjectDetectionModel(model, num_classes)\n",
    "\n",
    "# Get the number of output channels for class_logits and bbox_pred from the custom object detection model\n",
    "num_class_logits_channels = object_detection_model.cls_head[-1].out_features\n",
    "num_bbox_pred_channels = object_detection_model.bbox_head[-1].out_features\n",
    "\n",
    "# Print the number of output channels\n",
    "print(\"Number of output channels for class_logits:\", num_class_logits_channels)\n",
    "print(\"Number of output channels for bbox_pred:\", num_bbox_pred_channels)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "bea30c3e-d2e6-45fa-96ce-1c202c405e8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the loss functions\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "classification_criterion = nn.CrossEntropyLoss()\n",
    "regression_criterion = nn.SmoothL1Loss()\n",
    "\n",
    "learning_rate = 0.001 \n",
    "\n",
    "# Define the optimizer\n",
    "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "ceb9e5c3-e7e6-450a-8c09-0b88d5f5452b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1/3: 100%|██████████| 8490/8490 [26:23<00:00,  5.36it/s, loss=1.29e+3] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/3] Loss: 667.5331004577475\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2/3: 100%|██████████| 8490/8490 [18:22<00:00,  7.70it/s, loss=453]    \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [2/3] Loss: 651.836465227646\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 3/3: 100%|██████████| 8490/8490 [18:31<00:00,  7.64it/s, loss=705]    "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [3/3] Loss: 636.2471576672701\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm  # Import tqdm\n",
    "\n",
    "# Train the model\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    total_loss = 0.0\n",
    "    # Use tqdm to create a progress bar for the training loop\n",
    "    tqdm_data_loader = tqdm(train_data_loader, desc=f\"Epoch {epoch + 1}/{num_epochs}\")\n",
    "    for data_dict in tqdm_data_loader:\n",
    "        images = data_dict['images'].to(device)\n",
    "        class_labels = data_dict['class_labels'].to(device)\n",
    "        bboxes = data_dict['bboxes'].to(device)\n",
    "        widths = data_dict['widths'].to(device)\n",
    "        heights = data_dict['heights'].to(device)\n",
    "        \n",
    "        # Forward pass\n",
    "        cls_logits = model(images)[:8]\n",
    "        bbox_pred = model(images)[:8]\n",
    "\n",
    "        # class_logits, bbox_pred = model(images)\n",
    "        bbox_pred = bbox_pred[:, :4]\n",
    "\n",
    "        # Calculate classification loss\n",
    "        classification_loss = classification_criterion(cls_logits, class_labels)\n",
    "\n",
    "        # Calculate regression loss\n",
    "        regression_loss = regression_criterion(bbox_pred, bboxes)\n",
    "\n",
    "        # Calculate the total loss for this batch\n",
    "        batch_loss = classification_loss + regression_loss\n",
    "\n",
    "        # Backpropagation\n",
    "        optimizer.zero_grad()\n",
    "        batch_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # Accumulate the loss for this epoch\n",
    "        total_loss += batch_loss.item()\n",
    "\n",
    "        # Update the progress bar description with the current loss\n",
    "        tqdm_data_loader.set_postfix(loss=batch_loss.item())\n",
    "\n",
    "    # Print the average loss for the epoch\n",
    "    average_loss = total_loss / len(train_data_loader)\n",
    "    print(f\"Epoch [{epoch + 1}/{num_epochs}] Loss: {average_loss}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdec6b6e-df7e-4f17-b832-c8075f725028",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}