{
  "id": 405185,
  "title": "Is this model appropriate for bounding box regression?",
  "url": "/competitions/benetech-making-graphs-accessible/discussion/405185",
  "author_name": "Harry4463",
  "post_date": "2023-04-26T12:55:21.100000",
  "votes": 1,
  "comment_count": 0,
  "views": 0,
  "content": "<p>So I have this dataset: </p>\n<pre><code>class CustomDataset(Dataset):\n    def __init__(self, dataframe, image_dir, transform=None):\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.df[\"filename\"].unique())\n\n    def __getitem__(self, index):\n        image_id = self.df[\"filename\"].unique()[index]\n        image_path = os.path.join(self.image_dir, image_id)\n        image = Image.open(image_path).convert(\"RGB\")\n\n        #get the scale factor\n        scale_height = config.HEIGHT / image.size[1]\n        scale_width = config.WIDTH / image.size[0]\n\n\n        # filter the dataframe to only include rows with the current image_id\n        image_df = self.df[self.df[\"filename\"] == image_id]\n\n        boxes = []\n        labels = []\n\n        for i in range(len(image_df)):\n            # get the coordinates and label for the current box\n            x1 = image_df.iloc[i][\"x1\"] * scale_width\n            y1 = image_df.iloc[i][\"y1\"] * scale_height\n            x2 = image_df.iloc[i][\"x2\"] * scale_width\n            y2 = image_df.iloc[i][\"y2\"] * scale_height\n            x3 = image_df.iloc[i][\"x3\"] * scale_width\n            y3 = image_df.iloc[i][\"y3\"] * scale_height\n            x4 = image_df.iloc[i][\"x4\"] * scale_width\n            y4 = image_df.iloc[i][\"y4\"] * scale_height\n            label = image_df.iloc[i][\"label\"]\n\n            # add the box coordinates and label to the lists\n            boxes.append([x1, y1, x2, y2, x3, y3, x4, y4])\n            labels.append(label)\n\n        num_boxes = len(boxes)\n\n        # convert the boxes and labels to tensors\n        boxes = torch.tensor(boxes, dtype=torch.float32)\n        labels = torch.tensor(labels, dtype=torch.int64)\n\n        target = {}\n        target[\"image_id\"] = torch.tensor([index]).to(device)\n        target[\"boxes\"] = boxes.to(device)\n        target[\"labels\"] = labels.to(device)\n        target[\"num_boxes\"] = torch.tensor([num_boxes]).to(device)\n\n        if self.transform:\n            image = self.transform(image)\n\n        return image, target\n</code></pre>\n<p>And this model: </p>\n<pre><code>class RPN(nn.Module):\n    def __init__(self, num_bbox_params, num_classes):\n        super().__init__()\n        self.num_bbox_params = num_bbox_params\n        self.num_classes_params = num_classes\n        self.conv = nn.Conv2d(2048, 256, kernel_size=3, stride=1, padding=1).to(device)\n        self.cls_score = nn.Conv2d(256, self.num_classes_params, kernel_size=1, stride=1, padding=0)\n        self.bbox_pred = nn.Conv2d(256, self.num_bbox_params , kernel_size=1, stride=1, padding=0)\n\n    def forward(self, x, num_boxes):\n        features = self.conv(x).to(device)\n        num_bbox_params = num_boxes.sum().item() * self.num_bbox_params\n        num_classes_params = num_boxes.sum().item() * self.num_classes_params\n        self.cls_score = nn.Conv2d(256, num_classes_params, kernel_size=1, stride=1, padding=0).to(device)\n        cls_score = self.cls_score(features).to(device)\n        self.bbox_pred = nn.Conv2d(256, num_bbox_params, kernel_size=1, stride=1, padding=0).to(device)\n        bbox_pred = self.bbox_pred(features).to(device)\n        return cls_score, bbox_pred\n\nclass Regressor(nn.Module):\n    def __init__(self, num_bbox_params):\n        super().__init__()\n        self.num_bbox_params = num_bbox_params\n        self.regressor = nn.Sequential(\n            nn.Flatten(),\n            nn.Dropout(),\n            nn.Linear(2048, 512),\n            nn.ReLU(inplace=True),\n            nn.Linear(512, 128),\n            nn.ReLU(inplace=True),\n            nn.Linear(128, 32),\n            nn.ReLU(inplace=True),\n            nn.Linear(32, self.num_bbox_params)\n        ).to(device)\n\n    def forward(self, x, num_boxes):\n        # reshape features to (batch_size * num_boxes, num_features)\n        features = x.view(-1, 2048).to(device)\n        # dynamically set the output size of the final linear layer based on num_boxes\n        num_bbox_params = num_boxes.sum().item() * self.num_bbox_params\n        self.regressor[-1] = nn.Linear(32, num_bbox_params).to(device)\n        # apply regression to all bounding boxes in the batch\n        box = self.regressor(features).to(device)\n        # reshape box to (batch_size, num_boxes, num_bbox_params)\n        box = (box.view(-1, num_boxes.max().item(), self.num_bbox_params)).to(device)\n        return box\n\nclass Classifier(nn.Module):\n    def __init__(self, num_classes):\n        super().__init__()\n        self.num_classes_params = num_classes\n        self.classifier = nn.Sequential(\n            nn.Flatten(),\n            nn.Dropout(),\n            nn.Linear(2048, num_classes)\n        ).to(device)\n\n    def forward(self, x, num_boxes):\n        # reshape features to (batch_size * num_boxes, num_features)\n        features = x.view(-1, 2048).to(device)\n        # dynamically set the output size of the final linear layer based on num_boxes\n        num_classes_params = num_boxes.sum().item() * self.num_classes_params\n        self.classifier[-1] = nn.Linear(2048, num_classes_params).to(device)\n        # apply classification to all bounding boxes in the batch\n        classify = self.classifier(features).to(device)\n        # reshape classify to (batch_size, num_boxes, num_classes_params)\n        classify = (classify.view(-1, num_boxes.max().item(), self.num_classes_params)).to(device)\n        return classify\n\nclass BoundingBoxClassifier(nn.Module):\n    def __init__(self, num_classes, num_bbox_params):\n        super().__init__()\n        self.features = models.resnet50(weights = \"IMAGENET1K_V2\").to(device)\n        # remove the final classification layer of ResNet\n        self.features = nn.Sequential(*list(self.features.children())[:-1]).to(device)\n        self.rpn = RPN(num_bbox_params, num_classes).to(device)\n        self.regressor = Regressor(num_bbox_params).to(device)\n        self.classifier = Classifier(num_classes).to(device)\n\n\n    def forward(self, x, num_boxes):\n        features = self.features(x).to(device)\n        cls_score, bbox_pred = self.rpn(features, num_boxes)\n        box = self.regressor(features, num_boxes).to(device)\n        classify = self.classifier(features, num_boxes).to(device)\n        return box, classify, cls_score, bbox_pred\n</code></pre>\n<p>Just wondering how can I get the RPN head to work and play into regressor and classifier heads part of the model. Any help will be greatly appreciated</p>",
  "messages": [
    {
      "id": 2235961,
      "postDate": "2023-04-26T12:55:21.100Z",
      "content": "<p>So I have this dataset: </p>\n<pre><code>class CustomDataset(Dataset):\n    def __init__(self, dataframe, image_dir, transform=None):\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.df[\"filename\"].unique())\n\n    def __getitem__(self, index):\n        image_id = self.df[\"filename\"].unique()[index]\n        image_path = os.path.join(self.image_dir, image_id)\n        image = Image.open(image_path).convert(\"RGB\")\n\n        #get the scale factor\n        scale_height = config.HEIGHT / image.size[1]\n        scale_width = config.WIDTH / image.size[0]\n\n\n        # filter the dataframe to only include rows with the current image_id\n        image_df = self.df[self.df[\"filename\"] == image_id]\n\n        boxes = []\n        labels = []\n\n        for i in range(len(image_df)):\n            # get the coordinates and label for the current box\n            x1 = image_df.iloc[i][\"x1\"] * scale_width\n            y1 = image_df.iloc[i][\"y1\"] * scale_height\n            x2 = image_df.iloc[i][\"x2\"] * scale_width\n            y2 = image_df.iloc[i][\"y2\"] * scale_height\n            x3 = image_df.iloc[i][\"x3\"] * scale_width\n            y3 = image_df.iloc[i][\"y3\"] * scale_height\n            x4 = image_df.iloc[i][\"x4\"] * scale_width\n            y4 = image_df.iloc[i][\"y4\"] * scale_height\n            label = image_df.iloc[i][\"label\"]\n\n            # add the box coordinates and label to the lists\n            boxes.append([x1, y1, x2, y2, x3, y3, x4, y4])\n            labels.append(label)\n\n        num_boxes = len(boxes)\n\n        # convert the boxes and labels to tensors\n        boxes = torch.tensor(boxes, dtype=torch.float32)\n        labels = torch.tensor(labels, dtype=torch.int64)\n\n        target = {}\n        target[\"image_id\"] = torch.tensor([index]).to(device)\n        target[\"boxes\"] = boxes.to(device)\n        target[\"labels\"] = labels.to(device)\n        target[\"num_boxes\"] = torch.tensor([num_boxes]).to(device)\n\n        if self.transform:\n            image = self.transform(image)\n\n        return image, target\n</code></pre>\n<p>And this model: </p>\n<pre><code>class RPN(nn.Module):\n    def __init__(self, num_bbox_params, num_classes):\n        super().__init__()\n        self.num_bbox_params = num_bbox_params\n        self.num_classes_params = num_classes\n        self.conv = nn.Conv2d(2048, 256, kernel_size=3, stride=1, padding=1).to(device)\n        self.cls_score = nn.Conv2d(256, self.num_classes_params, kernel_size=1, stride=1, padding=0)\n        self.bbox_pred = nn.Conv2d(256, self.num_bbox_params , kernel_size=1, stride=1, padding=0)\n\n    def forward(self, x, num_boxes):\n        features = self.conv(x).to(device)\n        num_bbox_params = num_boxes.sum().item() * self.num_bbox_params\n        num_classes_params = num_boxes.sum().item() * self.num_classes_params\n        self.cls_score = nn.Conv2d(256, num_classes_params, kernel_size=1, stride=1, padding=0).to(device)\n        cls_score = self.cls_score(features).to(device)\n        self.bbox_pred = nn.Conv2d(256, num_bbox_params, kernel_size=1, stride=1, padding=0).to(device)\n        bbox_pred = self.bbox_pred(features).to(device)\n        return cls_score, bbox_pred\n\nclass Regressor(nn.Module):\n    def __init__(self, num_bbox_params):\n        super().__init__()\n        self.num_bbox_params = num_bbox_params\n        self.regressor = nn.Sequential(\n            nn.Flatten(),\n            nn.Dropout(),\n            nn.Linear(2048, 512),\n            nn.ReLU(inplace=True),\n            nn.Linear(512, 128),\n            nn.ReLU(inplace=True),\n            nn.Linear(128, 32),\n            nn.ReLU(inplace=True),\n            nn.Linear(32, self.num_bbox_params)\n        ).to(device)\n\n    def forward(self, x, num_boxes):\n        # reshape features to (batch_size * num_boxes, num_features)\n        features = x.view(-1, 2048).to(device)\n        # dynamically set the output size of the final linear layer based on num_boxes\n        num_bbox_params = num_boxes.sum().item() * self.num_bbox_params\n        self.regressor[-1] = nn.Linear(32, num_bbox_params).to(device)\n        # apply regression to all bounding boxes in the batch\n        box = self.regressor(features).to(device)\n        # reshape box to (batch_size, num_boxes, num_bbox_params)\n        box = (box.view(-1, num_boxes.max().item(), self.num_bbox_params)).to(device)\n        return box\n\nclass Classifier(nn.Module):\n    def __init__(self, num_classes):\n        super().__init__()\n        self.num_classes_params = num_classes\n        self.classifier = nn.Sequential(\n            nn.Flatten(),\n            nn.Dropout(),\n            nn.Linear(2048, num_classes)\n        ).to(device)\n\n    def forward(self, x, num_boxes):\n        # reshape features to (batch_size * num_boxes, num_features)\n        features = x.view(-1, 2048).to(device)\n        # dynamically set the output size of the final linear layer based on num_boxes\n        num_classes_params = num_boxes.sum().item() * self.num_classes_params\n        self.classifier[-1] = nn.Linear(2048, num_classes_params).to(device)\n        # apply classification to all bounding boxes in the batch\n        classify = self.classifier(features).to(device)\n        # reshape classify to (batch_size, num_boxes, num_classes_params)\n        classify = (classify.view(-1, num_boxes.max().item(), self.num_classes_params)).to(device)\n        return classify\n\nclass BoundingBoxClassifier(nn.Module):\n    def __init__(self, num_classes, num_bbox_params):\n        super().__init__()\n        self.features = models.resnet50(weights = \"IMAGENET1K_V2\").to(device)\n        # remove the final classification layer of ResNet\n        self.features = nn.Sequential(*list(self.features.children())[:-1]).to(device)\n        self.rpn = RPN(num_bbox_params, num_classes).to(device)\n        self.regressor = Regressor(num_bbox_params).to(device)\n        self.classifier = Classifier(num_classes).to(device)\n\n\n    def forward(self, x, num_boxes):\n        features = self.features(x).to(device)\n        cls_score, bbox_pred = self.rpn(features, num_boxes)\n        box = self.regressor(features, num_boxes).to(device)\n        classify = self.classifier(features, num_boxes).to(device)\n        return box, classify, cls_score, bbox_pred\n</code></pre>\n<p>Just wondering how can I get the RPN head to work and play into regressor and classifier heads part of the model. Any help will be greatly appreciated</p>",
      "rawMarkdown": "So I have this dataset: \n\n```\nclass CustomDataset(Dataset):\n    def __init__(self, dataframe, image_dir, transform=None):\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.df[\"filename\"].unique())\n\n    def __getitem__(self, index):\n        image_id = self.df[\"filename\"].unique()[index]\n        image_path = os.path.join(self.image_dir, image_id)\n        image = Image.open(image_path).convert(\"RGB\")\n\n        #get the scale factor\n        scale_height = config.HEIGHT / image.size[1]\n        scale_width = config.WIDTH / image.size[0]\n\n\n        # filter the dataframe to only include rows with the current image_id\n        image_df = self.df[self.df[\"filename\"] == image_id]\n\n        boxes = []\n        labels = []\n\n        for i in range(len(image_df)):\n            # get the coordinates and label for the current box\n            x1 = image_df.iloc[i][\"x1\"] * scale_width\n            y1 = image_df.iloc[i][\"y1\"] * scale_height\n            x2 = image_df.iloc[i][\"x2\"] * scale_width\n            y2 = image_df.iloc[i][\"y2\"] * scale_height\n            x3 = image_df.iloc[i][\"x3\"] * scale_width\n            y3 = image_df.iloc[i][\"y3\"] * scale_height\n            x4 = image_df.iloc[i][\"x4\"] * scale_width\n            y4 = image_df.iloc[i][\"y4\"] * scale_height\n            label = image_df.iloc[i][\"label\"]\n\n            # add the box coordinates and label to the lists\n            boxes.append([x1, y1, x2, y2, x3, y3, x4, y4])\n            labels.append(label)\n\n        num_boxes = len(boxes)\n\n        # convert the boxes and labels to tensors\n        boxes = torch.tensor(boxes, dtype=torch.float32)\n        labels = torch.tensor(labels, dtype=torch.int64)\n\n        target = {}\n        target[\"image_id\"] = torch.tensor([index]).to(device)\n        target[\"boxes\"] = boxes.to(device)\n        target[\"labels\"] = labels.to(device)\n        target[\"num_boxes\"] = torch.tensor([num_boxes]).to(device)\n\n        if self.transform:\n            image = self.transform(image)\n\n        return image, target\n```\n\n\nAnd this model: \n\n```\nclass RPN(nn.Module):\n    def __init__(self, num_bbox_params, num_classes):\n        super().__init__()\n        self.num_bbox_params = num_bbox_params\n        self.num_classes_params = num_classes\n        self.conv = nn.Conv2d(2048, 256, kernel_size=3, stride=1, padding=1).to(device)\n        self.cls_score = nn.Conv2d(256, self.num_classes_params, kernel_size=1, stride=1, padding=0)\n        self.bbox_pred = nn.Conv2d(256, self.num_bbox_params , kernel_size=1, stride=1, padding=0)\n\n    def forward(self, x, num_boxes):\n        features = self.conv(x).to(device)\n        num_bbox_params = num_boxes.sum().item() * self.num_bbox_params\n        num_classes_params = num_boxes.sum().item() * self.num_classes_params\n        self.cls_score = nn.Conv2d(256, num_classes_params, kernel_size=1, stride=1, padding=0).to(device)\n        cls_score = self.cls_score(features).to(device)\n        self.bbox_pred = nn.Conv2d(256, num_bbox_params, kernel_size=1, stride=1, padding=0).to(device)\n        bbox_pred = self.bbox_pred(features).to(device)\n        return cls_score, bbox_pred\n\nclass Regressor(nn.Module):\n    def __init__(self, num_bbox_params):\n        super().__init__()\n        self.num_bbox_params = num_bbox_params\n        self.regressor = nn.Sequential(\n            nn.Flatten(),\n            nn.Dropout(),\n            nn.Linear(2048, 512),\n            nn.ReLU(inplace=True),\n            nn.Linear(512, 128),\n            nn.ReLU(inplace=True),\n            nn.Linear(128, 32),\n            nn.ReLU(inplace=True),\n            nn.Linear(32, self.num_bbox_params)\n        ).to(device)\n\n    def forward(self, x, num_boxes):\n        # reshape features to (batch_size * num_boxes, num_features)\n        features = x.view(-1, 2048).to(device)\n        # dynamically set the output size of the final linear layer based on num_boxes\n        num_bbox_params = num_boxes.sum().item() * self.num_bbox_params\n        self.regressor[-1] = nn.Linear(32, num_bbox_params).to(device)\n        # apply regression to all bounding boxes in the batch\n        box = self.regressor(features).to(device)\n        # reshape box to (batch_size, num_boxes, num_bbox_params)\n        box = (box.view(-1, num_boxes.max().item(), self.num_bbox_params)).to(device)\n        return box\n\nclass Classifier(nn.Module):\n    def __init__(self, num_classes):\n        super().__init__()\n        self.num_classes_params = num_classes\n        self.classifier = nn.Sequential(\n            nn.Flatten(),\n            nn.Dropout(),\n            nn.Linear(2048, num_classes)\n        ).to(device)\n\n    def forward(self, x, num_boxes):\n        # reshape features to (batch_size * num_boxes, num_features)\n        features = x.view(-1, 2048).to(device)\n        # dynamically set the output size of the final linear layer based on num_boxes\n        num_classes_params = num_boxes.sum().item() * self.num_classes_params\n        self.classifier[-1] = nn.Linear(2048, num_classes_params).to(device)\n        # apply classification to all bounding boxes in the batch\n        classify = self.classifier(features).to(device)\n        # reshape classify to (batch_size, num_boxes, num_classes_params)\n        classify = (classify.view(-1, num_boxes.max().item(), self.num_classes_params)).to(device)\n        return classify\n\nclass BoundingBoxClassifier(nn.Module):\n    def __init__(self, num_classes, num_bbox_params):\n        super().__init__()\n        self.features = models.resnet50(weights = \"IMAGENET1K_V2\").to(device)\n        # remove the final classification layer of ResNet\n        self.features = nn.Sequential(*list(self.features.children())[:-1]).to(device)\n        self.rpn = RPN(num_bbox_params, num_classes).to(device)\n        self.regressor = Regressor(num_bbox_params).to(device)\n        self.classifier = Classifier(num_classes).to(device)\n        \n\n    def forward(self, x, num_boxes):\n        features = self.features(x).to(device)\n        cls_score, bbox_pred = self.rpn(features, num_boxes)\n        box = self.regressor(features, num_boxes).to(device)\n        classify = self.classifier(features, num_boxes).to(device)\n        return box, classify, cls_score, bbox_pred\n```\n\nJust wondering how can I get the RPN head to work and play into regressor and classifier heads part of the model. Any help will be greatly appreciated",
      "votes": 1
    }
  ],
  "comments": [],
  "raw_markdown_by_id": {
    "2235961": "So I have this dataset: \n\n```\nclass CustomDataset(Dataset):\n    def __init__(self, dataframe, image_dir, transform=None):\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.df[\"filename\"].unique())\n\n    def __getitem__(self, index):\n        image_id = self.df[\"filename\"].unique()[index]\n        image_path = os.path.join(self.image_dir, image_id)\n        image = Image.open(image_path).convert(\"RGB\")\n\n        #get the scale factor\n        scale_height = config.HEIGHT / image.size[1]\n        scale_width = config.WIDTH / image.size[0]\n\n\n        # filter the dataframe to only include rows with the current image_id\n        image_df = self.df[self.df[\"filename\"] == image_id]\n\n        boxes = []\n        labels = []\n\n        for i in range(len(image_df)):\n            # get the coordinates and label for the current box\n            x1 = image_df.iloc[i][\"x1\"] * scale_width\n            y1 = image_df.iloc[i][\"y1\"] * scale_height\n            x2 = image_df.iloc[i][\"x2\"] * scale_width\n            y2 = image_df.iloc[i][\"y2\"] * scale_height\n            x3 = image_df.iloc[i][\"x3\"] * scale_width\n            y3 = image_df.iloc[i][\"y3\"] * scale_height\n            x4 = image_df.iloc[i][\"x4\"] * scale_width\n            y4 = image_df.iloc[i][\"y4\"] * scale_height\n            label = image_df.iloc[i][\"label\"]\n\n            # add the box coordinates and label to the lists\n            boxes.append([x1, y1, x2, y2, x3, y3, x4, y4])\n            labels.append(label)\n\n        num_boxes = len(boxes)\n\n        # convert the boxes and labels to tensors\n        boxes = torch.tensor(boxes, dtype=torch.float32)\n        labels = torch.tensor(labels, dtype=torch.int64)\n\n        target = {}\n        target[\"image_id\"] = torch.tensor([index]).to(device)\n        target[\"boxes\"] = boxes.to(device)\n        target[\"labels\"] = labels.to(device)\n        target[\"num_boxes\"] = torch.tensor([num_boxes]).to(device)\n\n        if self.transform:\n            image = self.transform(image)\n\n        return image, target\n```\n\n\nAnd this model: \n\n```\nclass RPN(nn.Module):\n    def __init__(self, num_bbox_params, num_classes):\n        super().__init__()\n        self.num_bbox_params = num_bbox_params\n        self.num_classes_params = num_classes\n        self.conv = nn.Conv2d(2048, 256, kernel_size=3, stride=1, padding=1).to(device)\n        self.cls_score = nn.Conv2d(256, self.num_classes_params, kernel_size=1, stride=1, padding=0)\n        self.bbox_pred = nn.Conv2d(256, self.num_bbox_params , kernel_size=1, stride=1, padding=0)\n\n    def forward(self, x, num_boxes):\n        features = self.conv(x).to(device)\n        num_bbox_params = num_boxes.sum().item() * self.num_bbox_params\n        num_classes_params = num_boxes.sum().item() * self.num_classes_params\n        self.cls_score = nn.Conv2d(256, num_classes_params, kernel_size=1, stride=1, padding=0).to(device)\n        cls_score = self.cls_score(features).to(device)\n        self.bbox_pred = nn.Conv2d(256, num_bbox_params, kernel_size=1, stride=1, padding=0).to(device)\n        bbox_pred = self.bbox_pred(features).to(device)\n        return cls_score, bbox_pred\n\nclass Regressor(nn.Module):\n    def __init__(self, num_bbox_params):\n        super().__init__()\n        self.num_bbox_params = num_bbox_params\n        self.regressor = nn.Sequential(\n            nn.Flatten(),\n            nn.Dropout(),\n            nn.Linear(2048, 512),\n            nn.ReLU(inplace=True),\n            nn.Linear(512, 128),\n            nn.ReLU(inplace=True),\n            nn.Linear(128, 32),\n            nn.ReLU(inplace=True),\n            nn.Linear(32, self.num_bbox_params)\n        ).to(device)\n\n    def forward(self, x, num_boxes):\n        # reshape features to (batch_size * num_boxes, num_features)\n        features = x.view(-1, 2048).to(device)\n        # dynamically set the output size of the final linear layer based on num_boxes\n        num_bbox_params = num_boxes.sum().item() * self.num_bbox_params\n        self.regressor[-1] = nn.Linear(32, num_bbox_params).to(device)\n        # apply regression to all bounding boxes in the batch\n        box = self.regressor(features).to(device)\n        # reshape box to (batch_size, num_boxes, num_bbox_params)\n        box = (box.view(-1, num_boxes.max().item(), self.num_bbox_params)).to(device)\n        return box\n\nclass Classifier(nn.Module):\n    def __init__(self, num_classes):\n        super().__init__()\n        self.num_classes_params = num_classes\n        self.classifier = nn.Sequential(\n            nn.Flatten(),\n            nn.Dropout(),\n            nn.Linear(2048, num_classes)\n        ).to(device)\n\n    def forward(self, x, num_boxes):\n        # reshape features to (batch_size * num_boxes, num_features)\n        features = x.view(-1, 2048).to(device)\n        # dynamically set the output size of the final linear layer based on num_boxes\n        num_classes_params = num_boxes.sum().item() * self.num_classes_params\n        self.classifier[-1] = nn.Linear(2048, num_classes_params).to(device)\n        # apply classification to all bounding boxes in the batch\n        classify = self.classifier(features).to(device)\n        # reshape classify to (batch_size, num_boxes, num_classes_params)\n        classify = (classify.view(-1, num_boxes.max().item(), self.num_classes_params)).to(device)\n        return classify\n\nclass BoundingBoxClassifier(nn.Module):\n    def __init__(self, num_classes, num_bbox_params):\n        super().__init__()\n        self.features = models.resnet50(weights = \"IMAGENET1K_V2\").to(device)\n        # remove the final classification layer of ResNet\n        self.features = nn.Sequential(*list(self.features.children())[:-1]).to(device)\n        self.rpn = RPN(num_bbox_params, num_classes).to(device)\n        self.regressor = Regressor(num_bbox_params).to(device)\n        self.classifier = Classifier(num_classes).to(device)\n        \n\n    def forward(self, x, num_boxes):\n        features = self.features(x).to(device)\n        cls_score, bbox_pred = self.rpn(features, num_boxes)\n        box = self.regressor(features, num_boxes).to(device)\n        classify = self.classifier(features, num_boxes).to(device)\n        return box, classify, cls_score, bbox_pred\n```\n\nJust wondering how can I get the RPN head to work and play into regressor and classifier heads part of the model. Any help will be greatly appreciated"
  }
}