# -*- coding: utf-8 -*-

"""Untitled5.ipynb

Automatically generated by Colaboratory.

Original file is located at

    https://colab.research.google.com/drive/1nArNVlITuU8guqfbSRpu9a5IncIaPm_a

"""

from IPython.display import clear_output

!git clone https://github.com/matterport/Mask_RCNN.git # load Mask R-CNN code implementation

!git clone https://github.com/ruslan-kl/brain-tumor.git # load new data set and annotations

!pip install pycocotools

!rm -rf brain-tumor/.git/

!rm -rf Mask_RCNN/.git/

clear_output()

import os

import sys

from tqdm import tqdm

import cv2

import numpy as np

import json

import skimage.draw

import matplotlib

import matplotlib.pyplot as plt

import random

# Root directory of the project

ROOT_DIR = os.path.abspath('Mask_RCNN/')

# Import Mask RCNN

sys.path.append(ROOT_DIR)

from mrcnn.config import Config

from mrcnn import utils

from mrcnn.model import log

import mrcnn.model as modellib

from mrcnn import visualize

# Import COCO config

sys.path.append(os.path.join(ROOT_DIR, 'samples/coco/'))

import coco

plt.rcParams['figure.facecolor'] = 'white'

clear_output()

def get_ax(rows=1, cols=1, size=7):

    """Return a Matplotlib Axes array to be used in

    all visualizations in the notebook. Provide a

    central point to control graph sizes.

    Change the default size attribute to control the size

    of rendered images

    """

    _, ax = plt.subplots(rows, cols, figsize=(size*cols, size*rows))

    return ax

MODEL_DIR = os.path.join(ROOT_DIR, 'logs') # directory to save logs and trained model

# ANNOTATIONS_DIR = 'brain-tumor/data/new/annotations/' # directory with annotations for train/val sets

DATASET_DIR = 'brain-tumor/data_cleaned/' # directory with image data

DEFAULT_LOGS_DIR = 'logs'

# Local path to trained weights file

COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")

# Download COCO trained weights from Releases if needed

if not os.path.exists(COCO_MODEL_PATH):

    utils.download_trained_weights(COCO_MODEL_PATH)

class TumorConfig(Config):

    """Configuration for training on the brain tumor dataset.

    """

    # 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()

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)

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"]

)

# 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()

# 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.

print("Training network heads")

model.train(

    dataset_train, dataset_val,

    learning_rate=config.LEARNING_RATE,

    epochs=15,

    layers='heads'

)

# Recreate the model in inference mode

model = modellib.MaskRCNN(

    mode="inference",

    config=config,

    model_dir=DEFAULT_LOGS_DIR

)

# Get path to saved weights

# Either set a specific path or find last trained weights

# model_path = os.path.join(ROOT_DIR, ".h5 file name here")

model_path = model.find_last()

# Load trained weights

print("Loading weights from ", model_path)

model.load_weights(model_path, by_name=True)

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()

#vALIDATION SET

ind = 9

display_image(dataset_val, ind)

predict_and_plot_differences(dataset_val, ind)

ind = 6

display_image(dataset_val, ind)

predict_and_plot_differences(dataset_val, ind)

#Test Set

ind = 1

display_image(dataset_test, ind)

predict_and_plot_differences(dataset_test, ind)

ind = 0

display_image(dataset_test, ind)

predict_and_plot_differences(dataset_test, ind)