[ Object detection error ] AssertionError: The boxes tensor shape is not correct as Tensor[K, 5]

Teerath: <p>Hello everyone,</p>
<p>I am struggle to detect the object but facing the issue [AssertionError: The boxes tensor shape is not correct as Tensor[K, 5]]</p>
<p>Below is my code? Can you help me to fix?  I am not understanding.</p>
<p>import torch<br>
import torchvision<br>
from torchvision import transforms<br>
from torch.utils.data import DataLoader<br>
from torch.optim import SGD<br>
from torch.optim.lr_scheduler import MultiStepLR<br>
from torch.nn import functional as F<br>
from tqdm import tqdm<br>
import numpy as np<br>
import xml.etree.ElementTree as ET<br>
from sklearn.metrics import average_precision_score<br>
from torch.nn.utils.rnn import pad_sequence<br>
class_to_idx = {<br>
“background”: 0,<br>
“aeroplane”: 1,<br>
“bicycle”: 2,<br>
“bird”: 3,<br>
“boat”: 4,<br>
“bottle”: 5,<br>
“bus”: 6,<br>
“car”: 7,<br>
“cat”: 8,<br>
“chair”: 9,<br>
“cow”: 10,<br>
“diningtable”: 11,<br>
“dog”: 12,<br>
“horse”: 13,<br>
“motorbike”: 14,<br>
“person”: 15,<br>
“pottedplant”: 16,<br>
“sheep”: 17,<br>
“sofa”: 18,<br>
“train”: 19,<br>
“tvmonitor”: 20<br>
}<br>
def collate_fn_transform(batch):<br>
images = []<br>
targets = []<br>
for img, target in batch:<br>
images.append(img)</p>
<pre><code>    boxes_list = []
    labels_list = []
    for obj in target:
        boxes = obj['box']
        labels = obj['label']
        print(type(labels))
        boxes_list.append(torch.tensor(boxes, dtype=torch.float32))
        labels_list.append(torch.tensor(labels, dtype=torch.int64))
    
    boxes = torch.stack(boxes_list)
    labels = torch.stack(labels_list)

    targets.append({
        'boxes': boxes,
        'labels': labels
    })

return images, targets
</code></pre>
<h1>
<a name="def-collate_fn_transformbatch-1" class="anchor" href="#def-collate_fn_transformbatch-1"></a>def collate_fn_transform(batch):</h1>
<h1>
<a name="class_to_idx-2" class="anchor" href="#class_to_idx-2"></a>class_to_idx  = {</h1>
<h1>
<a name="background-0-3" class="anchor" href="#background-0-3"></a>“background”: 0,</h1>
<h1>
<a name="aeroplane-1-4" class="anchor" href="#aeroplane-1-4"></a>“aeroplane”: 1,</h1>
<h1>
<a name="bicycle-2-5" class="anchor" href="#bicycle-2-5"></a>“bicycle”: 2,</h1>
<h1>
<a name="bird-3-6" class="anchor" href="#bird-3-6"></a>“bird”: 3,</h1>
<h1>
<a name="boat-4-7" class="anchor" href="#boat-4-7"></a>“boat”: 4,</h1>
<h1>
<a name="bottle-5-8" class="anchor" href="#bottle-5-8"></a>“bottle”: 5,</h1>
<h1>
<a name="bus-6-9" class="anchor" href="#bus-6-9"></a>“bus”: 6,</h1>
<h1>
<a name="car-7-10" class="anchor" href="#car-7-10"></a>“car”: 7,</h1>
<h1>
<a name="cat-8-11" class="anchor" href="#cat-8-11"></a>“cat”: 8,</h1>
<h1>
<a name="chair-9-12" class="anchor" href="#chair-9-12"></a>“chair”: 9,</h1>
<h1>
<a name="cow-10-13" class="anchor" href="#cow-10-13"></a>“cow”: 10,</h1>
<h1>
<a name="diningtable-11-14" class="anchor" href="#diningtable-11-14"></a>“diningtable”: 11,</h1>
<h1>
<a name="dog-12-15" class="anchor" href="#dog-12-15"></a>“dog”: 12,</h1>
<h1>
<a name="horse-13-16" class="anchor" href="#horse-13-16"></a>“horse”: 13,</h1>
<h1>
<a name="motorbike-14-17" class="anchor" href="#motorbike-14-17"></a>“motorbike”: 14,</h1>
<h1>
<a name="person-15-18" class="anchor" href="#person-15-18"></a>“person”: 15,</h1>
<h1>
<a name="pottedplant-16-19" class="anchor" href="#pottedplant-16-19"></a>“pottedplant”: 16,</h1>
<h1>
<a name="sheep-17-20" class="anchor" href="#sheep-17-20"></a>“sheep”: 17,</h1>
<h1>
<a name="sofa-18-21" class="anchor" href="#sofa-18-21"></a>“sofa”: 18,</h1>
<h1>
<a name="train-19-22" class="anchor" href="#train-19-22"></a>“train”: 19,</h1>
<h1>
<a name="tvmonitor-20-23" class="anchor" href="#tvmonitor-20-23"></a>“tvmonitor”: 20</h1>
<h1>
<a name="h-24" class="anchor" href="#h-24"></a>}</h1>
<h1>
<a name="images-25" class="anchor" href="#images-25"></a>images = []</h1>
<h1>
<a name="targets-26" class="anchor" href="#targets-26"></a>targets = []</h1>
<h1>
<a name="for-img-target-in-batch-27" class="anchor" href="#for-img-target-in-batch-27"></a>for img, target in batch:</h1>
<h1>
<a name="imagesappendimg-28" class="anchor" href="#imagesappendimg-28"></a>images.append(img)</h1>
<h1>
<a name="targetsappendboxes-torchtensortargetboxes-dtypetorchfloat32-29" class="anchor" href="#targetsappendboxes-torchtensortargetboxes-dtypetorchfloat32-29"></a>targets.append({‘boxes’: torch.tensor(target[‘boxes’], dtype=torch.float32),</h1>
<h1>
<a name="labels-torchtensortargetlabels-dtypetorchint64-30" class="anchor" href="#labels-torchtensortargetlabels-dtypetorchint64-30"></a>‘labels’: torch.tensor(target[‘labels’], dtype=torch.int64)})</h1>
<h1>
<a name="return-images-targets-31" class="anchor" href="#return-images-targets-31"></a>return images, targets</h1>
<p>class_to_idx  = {<br>
“background”: 0,<br>
“aeroplane”: 1,<br>
“bicycle”: 2,<br>
“bird”: 3,<br>
“boat”: 4,<br>
“bottle”: 5,<br>
“bus”: 6,<br>
“car”: 7,<br>
“cat”: 8,<br>
“chair”: 9,<br>
“cow”: 10,<br>
“diningtable”: 11,<br>
“dog”: 12,<br>
“horse”: 13,<br>
“motorbike”: 14,<br>
“person”: 15,<br>
“pottedplant”: 16,<br>
“sheep”: 17,<br>
“sofa”: 18,<br>
“train”: 19,<br>
“tvmonitor”: 20<br>
}</p>
<h1>
<a name="define-dataset-class-for-voc2007-detection-32" class="anchor" href="#define-dataset-class-for-voc2007-detection-32"></a>Define dataset class for VOC2007 detection</h1>
<p>class VOC2007Detection(torchvision.datasets.VOCDetection):<br>
“”“VOC 2007 detection dataset.”“”</p>
<pre><code>def __init__(self, root, transform=None, target_transform=None, transforms=None):
    super(VOC2007Detection, self).__init__(root, year='2007', image_set='trainval',
                                           download=True, transform=transform,
                                           target_transform=target_transform, transforms=transforms)
    
def __getitem__(self, index):
    img, target = super(VOC2007Detection, self).__getitem__(index)

    # Convert target to tensor
    objects = []
    for obj in target['annotation']['object']:
        if obj['difficult'] == '0':
            bbox = obj['bndbox']
            objects.append({'box': [float(bbox['xmin']), float(bbox['ymin']),
                                    float(bbox['xmax']), float(bbox['ymax'])],
                             'label': int(class_to_idx[obj['name']])})

    target = objects

    return img, target
</code></pre>
<h1>
<a name="define-data-augmentation-transforms-33" class="anchor" href="#define-data-augmentation-transforms-33"></a>Define data augmentation transforms</h1>
<p>train_transforms = transforms.Compose([<br>
transforms.Resize((224, 224)),<br>
transforms.RandomHorizontalFlip(),<br>
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),<br>
transforms.ToTensor(),<br>
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),<br>
])</p>
<h1>
<a name="create-dataset-and-data-loader-34" class="anchor" href="#create-dataset-and-data-loader-34"></a>Create dataset and data loader</h1>
<p>train_dataset = VOC2007Detection(root=‘./data/VOCdevkit’, transform=train_transforms)<br>
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=2, shuffle=True, collate_fn=collate_fn_transform)</p>
<h1>
<a name="define-fast-r-cnn-model-35" class="anchor" href="#define-fast-r-cnn-model-35"></a>Define Fast R-CNN model</h1>
<p>class FastRCNN(torch.nn.Module):<br>
def <strong>init</strong>(self, num_classes):<br>
super(FastRCNN, self).<strong>init</strong>()</p>
<pre><code>    # Load pre-trained VGG16 backbone
    self.backbone = torchvision.models.vgg16(pretrained=True).features

    # Replace VGG16 classifier with a ROI pooling layer
    self.classifier = torch.nn.Sequential(torch.nn.Linear(25088, 4096),
                                           torch.nn.ReLU(inplace=True),
                                           torch.nn.Dropout(),
                                           torch.nn.Linear(4096, 4096),
                                           torch.nn.ReLU(inplace=True),
                                           torch.nn.Dropout(),
                                           torch.nn.Linear(4096, num_classes))

    self.roi_pooling = torchvision.ops.RoIPool(output_size=(7, 7), spatial_scale=0.0625)

    self.bbox_reg = torch.nn.Linear(4096, num_classes * 4)

def forward(self, x, proposals):
    x = self.backbone(x)
    x = x.view(x.size(0), -1)
    x = self.classifier(x)

    # ROI pooling
    roi_feats = []
    for proposal in proposals:
        roi_feats.append(self.roi_pooling(x.unsqueeze(0), proposal))

    roi_feats = torch.cat(roi_feats, dim=0)

    # Bounding box regression
    bbox_reg = self.bbox_reg(roi_feats)
    bbox_reg = bbox_reg.view(-1, num_classes, 4)

    return x, roi_feats, bbox_reg
</code></pre>
<h1>
<a name="define-device-36" class="anchor" href="#define-device-36"></a>Define device</h1>
<p>device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)</p>
<h1>
<a name="define-fast-r-cnn-model-37" class="anchor" href="#define-fast-r-cnn-model-37"></a>Define Fast R-CNN model</h1>
<h1>
<a name="extract-class-names-from-dataset-38" class="anchor" href="#extract-class-names-from-dataset-38"></a>Extract class names from dataset</h1>
<h1>
<a name="class_names-train_datasetclass_names-39" class="anchor" href="#class_names-train_datasetclass_names-39"></a>class_names = train_dataset.class_names</h1>
<h1>
<a name="create-class-to-index-dictionary-40" class="anchor" href="#create-class-to-index-dictionary-40"></a>Create class-to-index dictionary</h1>
<p>class_to_idx  = {<br>
“background”: 0,<br>
“aeroplane”: 1,<br>
“bicycle”: 2,<br>
“bird”: 3,<br>
“boat”: 4,<br>
“bottle”: 5,<br>
“bus”: 6,<br>
“car”: 7,<br>
“cat”: 8,<br>
“chair”: 9,<br>
“cow”: 10,<br>
“diningtable”: 11,<br>
“dog”: 12,<br>
“horse”: 13,<br>
“motorbike”: 14,<br>
“person”: 15,<br>
“pottedplant”: 16,<br>
“sheep”: 17,<br>
“sofa”: 18,<br>
“train”: 19,<br>
“tvmonitor”: 20<br>
}</p>
<h1>
<a name="define-fast-r-cnn-model-41" class="anchor" href="#define-fast-r-cnn-model-41"></a>Define Fast R-CNN model</h1>
<p>num_classes = 21<br>
model = FastRCNN(num_classes=num_classes).to(device)</p>
<h1>
<a name="define-loss-function-and-optimizer-42" class="anchor" href="#define-loss-function-and-optimizer-42"></a>Define loss function and optimizer</h1>
<p>criterion_cls = torch.nn.CrossEntropyLoss()<br>
criterion_reg = torch.nn.SmoothL1Loss()<br>
optimizer = SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=0.0005)<br>
scheduler = MultiStepLR(optimizer, milestones=[8, 10], gamma=0.1)</p>
<h1>
<a name="train-the-model-43" class="anchor" href="#train-the-model-43"></a>Train the model</h1>
<p>num_epochs = 12</p>
<p>for epoch in range(num_epochs):<br>
# Train<br>
model.train()<br>
train_loss_cls = 0<br>
train_loss_reg = 0<br>
num_batches = 0<br>
for images, targets in tqdm(train_loader):<br>
# images = images.to(device)<br>
images = [image.to(device) for image in images]<br>
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]<br>
images = torch.stack(images)<br>
# Generate region proposals<br>
proposals = []<br>
for target in targets:<br>
# proposal = torchvision.ops.box.convert_boxes_to_roi(target[‘boxes’])<br>
proposal = torchvision.ops.boxes.box_convert(target[‘boxes’], ‘xyxy’, ‘xywh’)<br>
proposals.append(proposal)<br>
# proposals = torch.stack(proposals)<br>
proposals = pad_sequence(proposals, batch_first=True, padding_value=0)</p>
<pre><code>    # Forward pass
    optimizer.zero_grad()
    _, roi_feats, bbox_reg = model(images, proposals)

    # Compute classification loss
    labels = torch.cat([t['labels'] for t in targets], dim=0)
    cls_scores = torch.cat([model.classifier(roi_feat.unsqueeze(0)) for roi_feat in roi_feats])
    loss_cls = criterion_cls(cls_scores, labels)

    # Compute bounding box regression loss
    gt_bbox_reg = torch.cat([t['boxes'] for t in targets], dim=0)
    gt_roi = torchvision.ops.box.convert_boxes_to_roi(gt_bbox_reg)
    mask = (labels &gt; 0).float()
    pred_bbox_reg = model.bbox_regressor(bbox_reg, proposals)
    loss_reg = criterion_reg(pred_bbox_reg * mask.unsqueeze(-1), gt_roi * mask.unsqueeze(-1))

    # Compute total loss
    loss = loss_cls + loss_reg
    train_loss_cls += loss_cls.item()
    train_loss_reg += loss_reg.item()
    num_batches += 1

    # Backward pass and optimization
    loss.backward()
    optimizer.step()
print("Done ", epoch)
</code></pre>