Faster R-CNN Over-Detecting

username: sahajp0
cooked: <p>Hi, I’m having trouble creating a Faster R-CNN around this dataset: <a href="https://figshare.com/articles/dataset/Annotated_dataset_for_deep-learning-based_bacterial_colony_detection/22022540/3" class="inline-onebox" rel="noopener nofollow ugc">Annotated dataset for deep-learning-based bacterial colony detection</a></p>
<p>It’s already in COCO labeled format, and I attempted to create an R-CNN using the following code (granted this is my first time with the architecture). However, the test image I feed into the model at the end always has MANY more bounding boxes than there are objects to detect. These boxes just don’t make sense either. The original paper ran the model at 100 epochs, so perhaps that is my issue, however, am I wrong to believe that the model should still somewhat accurately predict with only 10? Also, if I was to run 100 epochs, would you have any suggestions on how to feasibly run such a task. Since I’m using an M4 Pro Macbook, I don’t believe my hardware could really handle this much computation. Additionally, I would love to hear any feedback on ways that I could make my implementation faster, more accurate, or just in general more organized. Thank you!</p>
<p>model.py</p>
<pre><code class="lang-auto">import torch
from torch.utils.data import DataLoader
import torchvision
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.datasets import ImageFolder
from torchvision import transforms
import torchvision.transforms as T
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from datatset import CFUDataset
from dataset2 import CFUDataset2
import cv2
from PIL import Image

def test_accuracy(net, device="cpu"):
    correct = 0
    total = 0
    with torch.no_grad():
        for images, targets in valid_loader:
            images = list(image.to(device) for image in images)
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
            
            loss_dict = model(images, targets)
            print(loss_dict)

model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(weights="FasterRCNN_MobileNet_V3_Large_FPN_Weights.DEFAULT")

num_classes = 2 

in_features = model.roi_heads.box_predictor.cls_score.in_features

model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

transform = T.Compose([
    T.ToTensor(),
])

dataset = CFUDataset(root_dir="CFU_Dataset", annotations_json="CFU_Dataset/annot_COCO.json", transforms=transform)

indices = torch.randperm(len(dataset)).tolist()
train_dataset = torch.utils.data.Subset(dataset, indices[:-50])

valid_dataset = torch.utils.data.Subset(dataset, indices[-50:])

train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, 
                                   collate_fn=lambda x: tuple(zip(*x)))
valid_loader = DataLoader(valid_dataset, batch_size=8, shuffle=False, 
                                    collate_fn=lambda x: tuple(zip(*x)))

device = torch.device('cpu')
if torch.cuda.is_available():
    device = torch.device('cuda')

model.to(device)


params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.00001, momentum=0.9, 
                                                   weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, 
                                                               gamma=0.1)
num_epochs = 5
for epoch in range(num_epochs):
    model.train()
    train_loss = 0.0
    total = len(train_loader)
    count = 1
    for images, targets in train_loader:
        print(f"{count/total*100:.1f} %, Loss: {train_loss / count}",end="\r")
        count += 1
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        optimizer.zero_grad()

        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())

        losses.backward()
        optimizer.step()
        train_loss += losses.item()

    lr_scheduler.step()
    print(f'Epoch: {epoch + 1}, Loss: {train_loss / len(train_loader)}')
print("Training complete!")


model.eval()


import cv2
import torchvision.transforms as T
import torch

transform = T.Compose([
    T.ToTensor(),    # Convert PIL Image to Tensor
])
img = Image.open('CFU_Dataset/sp01_img02.jpg')

img = transform(img)  # Apply the transformation
img = img.float()  # Ensure the tensor is float32

with torch.no_grad():
    prediction = model([img])
img = cv2.imread('CFU_Dataset/sp01_img02.jpg')

for i in range(len(prediction[0]['boxes'])):
    x1, x2, x3, x4 = map(int, prediction[0]['boxes'][i].tolist())
    print(prediction[0]['scores'][i])
    image = cv2.rectangle(img, (x1, x2), (x3, x4), (255, 0, 0), 1)

cv2.imshow('CFU_Dataset/sp01_img02.jpg', image)
cv2.waitKey()

</code></pre>
<p>datatset.py</p>
<pre><code class="lang-auto">import torch
import json
import os
from PIL import Image

class CFUDataset(torch.utils.data.Dataset):

    def __init__(self, annotations_json, root_dir, transforms):
        file = open(annotations_json,)
        self.annotations = json.load(file)
        self.root_dir = root_dir
        self.transforms = transforms

    def __getitem__(self, idx):

        img_path = os.path.join(self.root_dir, self.annotations['images'][idx].get('file_name'))
        img_id = self.annotations['images'][idx].get('id')
        img = Image.open(img_path)

        boxes_array = []

        for annotation in self.annotations['annotations']:
            if annotation.get('image_id') == img_id:
                
                bbox = annotation.get('bbox')
                bbox = [bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]]
                
                if bbox[2] &gt; bbox[0] and bbox[3] &gt; bbox[1]:
                    boxes_array.append(bbox)
    

        boxes = torch.FloatTensor(list(boxes_array))
        
        labels = torch.tensor([1] * len(boxes_array), dtype=torch.int64)

        target = {}
        target["boxes"] = boxes.clone().detach()
        target["labels"] = labels.clone().detach()

        
        if self.transforms is not None:
            img = self.transforms(img)

        return img, target
    
    def __len__(self):
        return len(self.annotations['images']) 
</code></pre>
<p>Example output (white colonies are the objects, only one class):<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/f/0f6cd26168480523e3ac14a8816b9487826020b4.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/0f6cd26168480523e3ac14a8816b9487826020b4" title="Screenshot 2024-12-18 at 7.46.02 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f6cd26168480523e3ac14a8816b9487826020b4_2_500x500.jpeg" alt="Screenshot 2024-12-18 at 7.46.02 PM" data-base62-sha1="2csjk7mCXHupb6qB583hyURIvCk" width="500" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f6cd26168480523e3ac14a8816b9487826020b4_2_500x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f6cd26168480523e3ac14a8816b9487826020b4_2_750x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f6cd26168480523e3ac14a8816b9487826020b4_2_1000x1000.jpeg 2x" data-dominant-color="973A3F"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-12-18 at 7.46.02 PM</span><span class="informations">1670×1668 127 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>