5. Faster R-CNN¶
In chapter 4, we built a medical mask detection model using RetinaNet, a one-stage detector model. In this chapter, we will detect medical masks with Faster R-CNN, a two-stage detector.
From chapters 5.1 to 5.3, we will load the data, divide it into training and test data, and define the dataset class based on the code introduced in chapters 2 and 3. In chapter 5.4, we will use the torchvision API to load the pretrained model. In chapter 5.5, we will train the model through transfer learning, and finally, we will make inferences based on the test dataset and evaluate the model’s performance in chapter 5.6.
Before we begin the experiment, we should note that Google Colab allocates random GPUs, so a memory shortage may occur depending on the GPU that has been allotted.
It is recommended that you confirm that the GPU has enough memory before beginning the experiment. If you reset the runtime, you can be assigned a new GPU.
import torch
if torch.cuda.is_available():
device = torch.device("cuda")
print('There are %d GPU(s) available.' % torch.cuda.device_count())
print('We will use the GPU:', torch.cuda.get_device_name(0))
else:
print('No GPU available, using the CPU instead.')
device = torch.device("cpu")
There are 1 GPU(s) available.
We will use the GPU: Tesla T4
5.1 Loading the Data¶
We will load the data using the code from chapter 2.1. The code below describes the procedure for downloading and extracting the FaceMaskDetection dataset using the PL_data_loader.py file in Tutorial-Book-Utils repo, which is on PseudoLab’s Github.
!git clone https://github.com/Pseudo-Lab/Tutorial-Book-Utils
!python Tutorial-Book-Utils/PL_data_loader.py --data FaceMaskDetection
!unzip -q Face\ Mask\ Detection.zip
Cloning into 'Tutorial-Book-Utils'...
remote: Enumerating objects: 18, done.
remote: Counting objects: 5% (1/18)
remote: Counting objects: 11% (2/18)
remote: Counting objects: 16% (3/18)
remote: Counting objects: 22% (4/18)
remote: Counting objects: 27% (5/18)
remote: Counting objects: 33% (6/18)
remote: Counting objects: 38% (7/18)
remote: Counting objects: 44% (8/18)
remote: Counting objects: 50% (9/18)
remote: Counting objects: 55% (10/18)
remote: Counting objects: 61% (11/18)
remote: Counting objects: 66% (12/18)
remote: Counting objects: 72% (13/18)
remote: Counting objects: 77% (14/18)
remote: Counting objects: 83% (15/18)
remote: Counting objects: 88% (16/18)
remote: Counting objects: 94% (17/18)
remote: Counting objects: 100% (18/18)
remote: Counting objects: 100% (18/18), done.
remote: Compressing objects: 6% (1/15)
remote: Compressing objects: 13% (2/15)
remote: Compressing objects: 20% (3/15)
remote: Compressing objects: 26% (4/15)
remote: Compressing objects: 33% (5/15)
remote: Compressing objects: 40% (6/15)
remote: Compressing objects: 46% (7/15)
remote: Compressing objects: 53% (8/15)
remote: Compressing objects: 60% (9/15)
remote: Compressing objects: 66% (10/15)
remote: Compressing objects: 73% (11/15)
remote: Compressing objects: 80% (12/15)
remote: Compressing objects: 86% (13/15)
remote: Compressing objects: 93% (14/15)
remote: Compressing objects: 100% (15/15)
remote: Compressing objects: 100% (15/15), done.
remote: Total 18 (delta 4), reused 8 (delta 2), pack-reused 0
Unpacking objects: 5% (1/18)
Unpacking objects: 11% (2/18)
Unpacking objects: 16% (3/18)
Unpacking objects: 22% (4/18)
Unpacking objects: 27% (5/18)
Unpacking objects: 33% (6/18)
Unpacking objects: 38% (7/18)
Unpacking objects: 44% (8/18)
Unpacking objects: 50% (9/18)
Unpacking objects: 55% (10/18)
Unpacking objects: 61% (11/18)
Unpacking objects: 66% (12/18)
Unpacking objects: 72% (13/18)
Unpacking objects: 77% (14/18)
Unpacking objects: 83% (15/18)
Unpacking objects: 88% (16/18)
Unpacking objects: 94% (17/18)
Unpacking objects: 100% (18/18)
Unpacking objects: 100% (18/18), done.
Face Mask Detection.zip is done!
5.2 Data Separation¶
We will separate the dataset as in chapter 3.3. By using the code below, 170 images are extracted randomly and moved to the test folder.
import os
import random
import numpy as np
import shutil
print(len(os.listdir('annotations')))
print(len(os.listdir('images')))
!mkdir test_images
!mkdir test_annotations
random.seed(1234)
idx = random.sample(range(853), 170)
for img in np.array(sorted(os.listdir('images')))[idx]:
shutil.move('images/'+img, 'test_images/'+img)
for annot in np.array(sorted(os.listdir('annotations')))[idx]:
shutil.move('annotations/'+annot, 'test_annotations/'+annot)
print(len(os.listdir('annotations')))
print(len(os.listdir('images')))
print(len(os.listdir('test_annotations')))
print(len(os.listdir('test_images')))
853
853
683
683
170
170
We will also load the packages needed for modeling. torchvision is used for image processing and has built-in packages for datasets and models.
import os
import numpy as np
import matplotlib.patches as patches
import matplotlib.pyplot as plt
from bs4 import BeautifulSoup
from PIL import Image
import torchvision
from torchvision import transforms, datasets, models
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
import time
Defining the Dataset Class¶
This time, we will define the functions for the bounding boxes as shown in chapter 2.3.
def generate_box(obj):
xmin = float(obj.find('xmin').text)
ymin = float(obj.find('ymin').text)
xmax = float(obj.find('xmax').text)
ymax = float(obj.find('ymax').text)
return [xmin, ymin, xmax, ymax]
adjust_label = 1
def generate_label(obj):
if obj.find('name').text == "with_mask":
return 1 + adjust_label
elif obj.find('name').text == "mask_weared_incorrect":
return 2 + adjust_label
return 0 + adjust_label
def generate_target(file):
with open(file) as f:
data = f.read()
soup = BeautifulSoup(data, "html.parser")
objects = soup.find_all("object")
num_objs = len(objects)
boxes = []
labels = []
for i in objects:
boxes.append(generate_box(i))
labels.append(generate_label(i))
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.as_tensor(labels, dtype=torch.int64)
target = {}
target["boxes"] = boxes
target["labels"] = labels
return target
def plot_image_from_output(img, annotation):
img = img.cpu().permute(1,2,0)
fig,ax = plt.subplots(1)
ax.imshow(img)
for idx in range(len(annotation["boxes"])):
xmin, ymin, xmax, ymax = annotation["boxes"][idx]
if annotation['labels'][idx] == 1 :
rect = patches.Rectangle((xmin,ymin),(xmax-xmin),(ymax-ymin),linewidth=1,edgecolor='r',facecolor='none')
elif annotation['labels'][idx] == 2 :
rect = patches.Rectangle((xmin,ymin),(xmax-xmin),(ymax-ymin),linewidth=1,edgecolor='g',facecolor='none')
else :
rect = patches.Rectangle((xmin,ymin),(xmax-xmin),(ymax-ymin),linewidth=1,edgecolor='orange',facecolor='none')
ax.add_patch(rect)
plt.show()
Next, as in chapter 4.3, we will define the dataset class and data loader. The dataset will be loaded and the batch size set to 4 through the `torch.utils.data.DataLoader{/code0} function. You can change the batch size according to your individual memory size.
class MaskDataset(object):
def __init__(self, transforms, path):
'''
path: path to train folder or test folder
'''
# define the path to the images and what transform will be used
self.transforms = transforms
self.path = path
self.imgs = list(sorted(os.listdir(self.path)))
def __getitem__(self, idx): #special method
# load images ad masks
file_image = self.imgs[idx]
file_label = self.imgs[idx][:-3] + 'xml'
img_path = os.path.join(self.path, file_image)
if 'test' in self.path:
label_path = os.path.join("test_annotations/", file_label)
else:
label_path = os.path.join("annotations/", file_label)
img = Image.open(img_path).convert("RGB")
#Generate Label
target = generate_target(label_path)
if self.transforms is not None:
img = self.transforms(img)
return img, target
def __len__(self):
return len(self.imgs)
data_transform = transforms.Compose([ # transforms.Compose : a class that calls the functions in a list consecutively
transforms.ToTensor() # ToTensor : convert numpy image to torch.Tensor type
])
def collate_fn(batch):
return tuple(zip(*batch))
dataset = MaskDataset(data_transform, 'images/')
test_dataset = MaskDataset(data_transform, 'test_images/')
data_loader = torch.utils.data.DataLoader(dataset, batch_size=4, collate_fn=collate_fn)
test_data_loader = torch.utils.data.DataLoader(test_dataset, batch_size=2, collate_fn=collate_fn)
5.4 Import Model¶
torchvision.models.detection provides the Faster R-CNN API (torchvision.models.detection.fasterrcnn_resnet50_fpn) so it can be easily implemented. This provides a model that has been pre-trained with the COCO dataset using ResNet50. We can choose to load the pre-trained weights by declaring pretrained=True/False.
When loading the model, set the desired number of classes in num_classes and use the model. One thing to note when using Faster R-CNN is that you should include the background class when specifying the class number in num_classes. In other words, you need to increase the number of classes in the dataset by one in order to add the background class.
def get_model_instance_segmentation(num_classes):
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
5.5 Transfer Learning¶
We will now perform transfer learning on the Face Mask Detection dataset. The Face Mask Detection dataset consists of 3 classes, so we will load the model by setting num_classes to 4, making sure to include the background class.
If the current environment allows it, send the model to the GPU memory in order to speed up training.
model = get_model_instance_segmentation(4)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)
FasterRCNN(
(transform): GeneralizedRCNNTransform(
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
Resize(min_size=(800,), max_size=1333, mode='bilinear')
)
(backbone): BackboneWithFPN(
(body): IntermediateLayerGetter(
(conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
(bn1): FrozenBatchNorm2d(64)
(relu): ReLU(inplace=True)
(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(layer1): Sequential(
(0): Bottleneck(
(conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(64)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(64)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(256)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(1): FrozenBatchNorm2d(256)
)
)
(1): Bottleneck(
(conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(64)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(64)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(256)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(64)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(64)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(256)
(relu): ReLU(inplace=True)
)
)
(layer2): Sequential(
(0): Bottleneck(
(conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): FrozenBatchNorm2d(512)
)
)
(1): Bottleneck(
(conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512)
(relu): ReLU(inplace=True)
)
(3): Bottleneck(
(conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512)
(relu): ReLU(inplace=True)
)
)
(layer3): Sequential(
(0): Bottleneck(
(conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): FrozenBatchNorm2d(1024)
)
)
(1): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024)
(relu): ReLU(inplace=True)
)
(3): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024)
(relu): ReLU(inplace=True)
)
(4): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024)
(relu): ReLU(inplace=True)
)
(5): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024)
(relu): ReLU(inplace=True)
)
)
(layer4): Sequential(
(0): Bottleneck(
(conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(512)
(conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(512)
(conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(2048)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): FrozenBatchNorm2d(2048)
)
)
(1): Bottleneck(
(conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(512)
(conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(512)
(conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(2048)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(512)
(conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(512)
(conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(2048)
(relu): ReLU(inplace=True)
)
)
)
(fpn): FeaturePyramidNetwork(
(inner_blocks): ModuleList(
(0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))
(1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))
(2): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))
(3): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))
)
(layer_blocks): ModuleList(
(0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
)
(extra_blocks): LastLevelMaxPool()
)
)
(rpn): RegionProposalNetwork(
(anchor_generator): AnchorGenerator()
(head): RPNHead(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(cls_logits): Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))
(bbox_pred): Conv2d(256, 12, kernel_size=(1, 1), stride=(1, 1))
)
)
(roi_heads): RoIHeads(
(box_roi_pool): MultiScaleRoIAlign()
(box_head): TwoMLPHead(
(fc6): Linear(in_features=12544, out_features=1024, bias=True)
(fc7): Linear(in_features=1024, out_features=1024, bias=True)
)
(box_predictor): FastRCNNPredictor(
(cls_score): Linear(in_features=1024, out_features=4, bias=True)
(bbox_pred): Linear(in_features=1024, out_features=16, bias=True)
)
)
)
You can see what layers Faster R-CNN is composed of through the above output. Also, the availability of the GPU can be seen through torch.cuda.is_available().
torch.cuda.is_available()
True
Now that the model has been initiated, we will train it. The parameter for setting the number of epochs for training is (num_epochs) and will be set to 10. We will optimize the model using the SGD method. Each hyper parameter can be modified for the user’s needs.
num_epochs = 10
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,
momentum=0.9, weight_decay=0.0005)
Now we will train the model. Using the data_loader created above, we will input the data into the model in batches. Then, we will calculate the loss and optimize the model weights. By observing the loss printed at every epoch, we can see that the training process is being executed.
print('----------------------train start--------------------------')
for epoch in range(num_epochs):
start = time.time()
model.train()
i = 0
epoch_loss = 0
for imgs, annotations in data_loader:
i += 1
imgs = list(img.to(device) for img in imgs)
annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
loss_dict = model(imgs, annotations)
losses = sum(loss for loss in loss_dict.values())
optimizer.zero_grad()
losses.backward()
optimizer.step()
epoch_loss += losses
print(f'epoch : {epoch+1}, Loss : {epoch_loss}, time : {time.time() - start}')
----------------------train start--------------------------
epoch : 1, Loss : 77.14759063720703, time : 252.42370867729187
epoch : 2, Loss : 48.91315460205078, time : 263.22984743118286
epoch : 3, Loss : 43.18947982788086, time : 264.4591932296753
epoch : 4, Loss : 36.07373046875, time : 265.2568733692169
epoch : 5, Loss : 31.8864688873291, time : 265.57766008377075
epoch : 6, Loss : 31.76308250427246, time : 265.0076003074646
epoch : 7, Loss : 31.24744415283203, time : 265.16882514953613
epoch : 8, Loss : 29.340274810791016, time : 265.73448038101196
epoch : 9, Loss : 25.922008514404297, time : 267.91367626190186
epoch : 10, Loss : 23.59230613708496, time : 266.9004054069519
If we want to save the trained weights, we can use torch.save to save and use the code underneath it to load it when needed.
torch.save(model.state_dict(),f'model_{num_epochs}.pt')
model.load_state_dict(torch.load(f'model_{num_epochs}.pt'))
<All keys matched successfully>
5.6 Inference¶
Since the model has been trained, we will check the inference results to see if it has been trained well. The prediction results include the bounding box coordinates (boxes), classes (labels), and confidence scores (scores). Under the confidence scores (scores), the confidence value of the corresponding class is stored. We will define the function make_prediction to extract only those with a threshold of 0.5 or higher. Then we will print the results of the first batch of the test_data_loader.
def make_prediction(model, img, threshold):
model.eval()
preds = model(img)
for id in range(len(preds)) :
idx_list = []
for idx, score in enumerate(preds[id]['scores']) :
if score > threshold :
idx_list.append(idx)
preds[id]['boxes'] = preds[id]['boxes'][idx_list]
preds[id]['labels'] = preds[id]['labels'][idx_list]
preds[id]['scores'] = preds[id]['scores'][idx_list]
return preds
with torch.no_grad():
# batch size of the test set = 2
for imgs, annotations in test_data_loader:
imgs = list(img.to(device) for img in imgs)
pred = make_prediction(model, imgs, 0.5)
print(pred)
break
[{'boxes': tensor([[117.7811, 1.4936, 132.9596, 18.4192],
[214.8204, 59.8669, 249.7893, 97.6275]], device='cuda:0'), 'labels': tensor([2, 2], device='cuda:0'), 'scores': tensor([0.9430, 0.9414], device='cuda:0')}, {'boxes': tensor([[218.8598, 99.3362, 260.0332, 138.8516],
[130.5172, 109.1189, 179.2908, 152.5566],
[ 29.2499, 88.7732, 45.5664, 104.5635],
[ 40.9168, 109.1093, 67.3653, 140.0567],
[165.5889, 90.0294, 179.4471, 109.1606],
[ 83.7276, 84.3918, 94.5928, 96.4693],
[302.4648, 130.4534, 332.0580, 158.8674],
[258.4624, 90.7134, 269.2498, 102.2883],
[ 2.8419, 103.6409, 21.9580, 125.5492]], device='cuda:0'), 'labels': tensor([2, 2, 1, 1, 1, 1, 1, 1, 1], device='cuda:0'), 'scores': tensor([0.9962, 0.9918, 0.9900, 0.9894, 0.9891, 0.9653, 0.9652, 0.9573, 0.9046],
device='cuda:0')}]
Using the predicted result, we will draw the bounding boxes on the images. The image is plotted using the plot_image_from_output function defined above. Target is the actual position of the bounding boxes and Prediction is the predicted result from the model. We can see that the model has located the actual position of the bounding boxes well.
_idx = 1
print("Target : ", annotations[_idx]['labels'])
plot_image_from_output(imgs[_idx], annotations[_idx])
print("Prediction : ", pred[_idx]['labels'])
plot_image_from_output(imgs[_idx], pred[_idx])
Target : tensor([1, 1, 1, 2, 2, 1, 1, 1])
Prediction : tensor([2, 2, 1, 1, 1, 1, 1, 1, 1], device='cuda:0')
This time, we will evaluate the inference results for all of the test data. First, the predicted results and actual labels for all test data will be saved in preds_adj_all and annot_all, respectively.
from tqdm import tqdm
labels = []
preds_adj_all = []
annot_all = []
for im, annot in tqdm(test_data_loader, position = 0, leave = True):
im = list(img.to(device) for img in im)
for t in annot:
labels += t['labels']
with torch.no_grad():
preds_adj = make_prediction(model, im, 0.5)
preds_adj = [{k: v.to(torch.device('cpu')) for k, v in t.items()} for t in preds_adj]
preds_adj_all.append(preds_adj)
annot_all.append(annot)
100%|██████████| 85/85 [00:25<00:00, 3.34it/s]
Then, we will use the utils_ObjectDetection.py file in the Tutorial-Book-Utils folder to calculate the mAP value. the get_batch_statistics function is utilized to calculate several statistics using only bounding boxes with an IoU (Intersection over Union) value higher than 0.5. After that, the ap_per_class function is used to calculate the AP value for each class.
%cd Tutorial-Book-Utils/
import utils_ObjectDetection as utils
/content/Tutorial-Book-Utils
sample_metrics = []
for batch_i in range(len(preds_adj_all)):
sample_metrics += utils.get_batch_statistics(preds_adj_all[batch_i], annot_all[batch_i], iou_threshold=0.5)
true_positives, pred_scores, pred_labels = [torch.cat(x, 0) for x in list(zip(*sample_metrics))] # all the batches get concatenated
precision, recall, AP, f1, ap_class = utils.ap_per_class(true_positives, pred_scores, pred_labels, torch.tensor(labels))
mAP = torch.mean(AP)
print(f'mAP : {mAP}')
print(f'AP : {AP}')
mAP : 0.7182363990382057
AP : tensor([0.8694, 0.9189, 0.3664], dtype=torch.float64)
AP values are shown only for the 3 actual classes, excluding the background class. Even after training only 10 epochs, you can observe that it is better than the RetinaNet results from chapter 4. The model has a notable AP value of 0.9189 when detecting objects wearing a mask, which is annotated with class 1. It also shows an AP of 0.3664 for class 2, which represents objects that are not properly wearing a mask. It is well known that RetinaNet shows high performance despite being a one-stage method with FPN and focal loss, but it seems it is not suitable for this particular dataset. It may be also due to the lack of hyperparameter tuning.
This concludes the medical mask detection tutorial. Throughout this tutorial, we went from preprocessing the dataset to training and predicting the model. To achieve better results, you can increase the number of epochs for training or try hyperparameter tuning. Try applying the object detection model on the data you want and evaluate the results.