Brain Tumor Detection using Mask R-CNN

Original article can be found here (source): Deep Learning on Medium

Step-3: Configuration for training on the brain tumor dataset.

Here we need to setup configuration include properties like setting the number of GPUs to use along with the number of images per GPU, Number of classes (we would normally add +1 for the background), Number of training steps per epoch, Learning rate, Skip detections with < 85% confidence,

class TumorConfig(Config):
# Give the configuration a recognizable name
NAME = 'tumor_detector'
GPU_COUNT = 1
IMAGES_PER_GPU = 1
NUM_CLASSES = 1 + 1 # background + tumor
DETECTION_MIN_CONFIDENCE = 0.85
STEPS_PER_EPOCH = 100
LEARNING_RATE = 0.001

config = TumorConfig()
config.display()

Step: 4: Build the custom brain MRI data set.

Dataset class provides a consistent way to work with any dataset. We will create our new datasets for brain images to train without having to change the code of the model.

Dataset class also supports loading multiple data sets at the same time. This is very helpful when you want to detect different objects and they are all not available in one data set.

In the load_dataset method, we iterate through all the files in the image and annotations folders to add the class, images, and annotations to create the dataset using add_class and add_image methods.

extract_boxes method extracts each of the bounding boxes from the annotation file. Annotation files are XML files using Pascal VOC format. It returns the box, it’s height and width

load_mask method generates the masks for every object in the image. It returns one mask per instance and class ids, a 1D array of class id for the instance masks

image_reference method returns the path of the image.

class BrainScanDataset(utils.Dataset):def load_brain_scan(self, dataset_dir, subset):
"""Load a subset of the FarmCow dataset.
dataset_dir: Root directory of the dataset.
subset: Subset to load: train or val
"""
# Add classes. We have only one class to add.
self.add_class("tumor", 1, "tumor")
# Train or validation dataset?
assert subset in ["train", "val", 'test']
dataset_dir = os.path.join(dataset_dir, subset)
annotations = json.load(open(os.path.join(DATASET_DIR, subset, 'annotations_'+subset+'.json')))
annotations = list(annotations.values()) # don't need the dict keys
# The VIA tool saves images in the JSON even if they don't have any
# annotations. Skip unannotated images.
annotations = [a for a in annotations if a['regions']]
# Add images
for a in annotations:
# Get the x, y coordinaets of points of the polygons that make up
# the outline of each object instance. These are stores in the
# shape_attributes (see json format above)
# The if condition is needed to support VIA versions 1.x and 2.x.
if type(a['regions']) is dict:
polygons = [r['shape_attributes'] for r in a['regions'].values()]
else:
polygons = [r['shape_attributes'] for r in a['regions']]
# load_mask() needs the image size to convert polygons to masks.
# Unfortunately, VIA doesn't include it in JSON, so we must read
# the image. This is only managable since the dataset is tiny.
image_path = os.path.join(dataset_dir, a['filename'])
image = skimage.io.imread(image_path)
height, width = image.shape[:2]
self.add_image(
"tumor",
image_id=a['filename'], # use file name as a unique image id
path=image_path,
width=width,
height=height,
polygons=polygons
)
def load_mask(self, image_id):
"""Generate instance masks for an image.
Returns:
masks: A bool array of shape [height, width, instance count] with
one mask per instance.
class_ids: a 1D array of class IDs of the instance masks.
"""
# If not a farm_cow dataset image, delegate to parent class.
image_info = self.image_info[image_id]
if image_info["source"] != "tumor":
return super(self.__class__, self).load_mask(image_id)
# Convert polygons to a bitmap mask of shape
# [height, width, instance_count]
info = self.image_info[image_id]
mask = np.zeros([info["height"], info["width"], len(info["polygons"])],
dtype=np.uint8)
for i, p in enumerate(info["polygons"]):
# Get indexes of pixels inside the polygon and set them to 1
rr, cc = skimage.draw.polygon(p['all_points_y'], p['all_points_x'])
mask[rr, cc, i] = 1
# Return mask, and array of class IDs of each instance. Since we have
# one class ID only, we return an array of 1s
return mask.astype(np.bool), np.ones([mask.shape[-1]], dtype=np.int32)
def image_reference(self, image_id):
"""Return the path of the image."""
info = self.image_info[image_id]
if info["source"] == "tumor":
return info["path"]
else:
super(self.__class__, self).image_reference(image_id)

Step-5: Initialize the Mask R-CNN model for training using the Config instance that we created and load the pre-trained weights for the Mask R-CNN from the COCO data set excluding the last few layers.

Since we’re using a very small dataset, and starting from COCO trained weights, we don’t need to train too long. Also, no need to train all layers, just the heads should do it.

We exclude the last few layers from training for ResNet101.

model = modellib.MaskRCNN(
mode='training',
config=config,
model_dir=DEFAULT_LOGS_DIR
)
model.load_weights(
COCO_MODEL_PATH,
by_name=True,
exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"]
)

Step: 6: Load the dataset and train your model for 15 epochs with the Learning rate as 0.001

# Training dataset.
dataset_train = BrainScanDataset()
dataset_train.load_brain_scan(DATASET_DIR, 'train')
dataset_train.prepare()
# Validation dataset
dataset_val = BrainScanDataset()
dataset_val.load_brain_scan(DATASET_DIR, 'val')
dataset_val.prepare()
dataset_test = BrainScanDataset()
dataset_test.load_brain_scan(DATASET_DIR, 'test')
dataset_test.prepare()
print("Training network heads")
model.train(
dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE,
epochs=15,
layers='heads'
)

Step 7: Recreate the model in inference mode

# Recreate the model in inference mode
model = modellib.MaskRCNN(
mode="inference",
config=config,
model_dir=DEFAULT_LOGS_DIR
)
model_path = model.find_last()# Load trained weights
print("Loading weights from ", model_path)
model.load_weights(model_path, by_name=True)

Step 8: Now build functions to display the results.

def predict_and_plot_differences(dataset, img_id):
original_image, image_meta, gt_class_id, gt_box, gt_mask =\
modellib.load_image_gt(dataset, config,
img_id, use_mini_mask=False)
results = model.detect([original_image], verbose=0)
r = results[0]
visualize.display_differences(
original_image,
gt_box, gt_class_id, gt_mask,
r['rois'], r['class_ids'], r['scores'], r['masks'],
class_names = ['tumor'], title="", ax=get_ax(),
show_mask=True, show_box=True)
def display_image(dataset, ind):
plt.figure(figsize=(5,5))
plt.imshow(dataset.load_image(ind))
plt.xticks([])
plt.yticks([])
plt.title('Original Image')
plt.show()

Test your model prediction on the validation set:

#Validation set
ind = 9
display_image(dataset_val, ind)
predict_and_plot_differences(dataset_val, ind)