projfinal_maskrcnn_basemodel.py

# -*- coding: utf-8 -*-
"""ProjFinal_MaskRCNN_baseModel.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1LRq_rDaCmJydu41lGMKDb7UokJD72oIW
"""

!pip install keras==2.2.5

# Commented out IPython magic to ensure Python compatibility.
# load built-in libraries
import os
import sys
import random
import math
import re
import time
import pandas as pd
import numpy as np 
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
# %tensorflow_version 1.x
import tensorflow as tf
import keras

from google.colab import drive
# %matplotlib inline
from skimage import io, color
from pycocotools.coco import COCO

# mount drive on Google drive to access training data
# Ignore this if you don't use Google Colab
drive.mount("/content/drive")

# import Mask-RCNN
rcnn_path = "/content/drive/My Drive/ECE_542/Project-final/scripts/Mask_RCNN_master/"
sys.path.append(rcnn_path)

from mrcnn.config import Config
import mrcnn.utils as utils
import mrcnn.model as modellib
import mrcnn.visualize as visualize
from mrcnn.model import log

"""# Access important directories"""

ROOT_DIR = "/content/drive/My Drive/ECE_542/Project-final/"
MODEL_DIR = os.path.join(ROOT_DIR, "logs/")
COCO_MODEL_PATH = os.path.join(rcnn_path, "mrcnn/mask_rcnn_coco.h5")
IMAGE_DIR = os.path.join(ROOT_DIR, "data_resized/")
TRAIN_PATH = os.path.join(ROOT_DIR, "train_resize.json")
VAL_PATH = os.path.join(ROOT_DIR, "test_resize.json")

"""# Model configurations"""

# Configuration of base model
class base_modelConfig(Config):  #
    """Configuration for training on the cigarette butts dataset.
    Derives from the base Config class and overrides values specific
    to the cigarette butts dataset.
    """
    # Give the configuration a recognizable name
    NAME = "plant_phenotyping"

    # Train on 1 GPU and 1 image per GPU. Batch size is 1 (GPUs * images/GPU).
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

    # Number of classes (including background)
    NUM_CLASSES = 1 + 2  # background + 1 (cig_butt)

    # All of our training images are 512x512
    IMAGE_MIN_DIM = 512
    IMAGE_MAX_DIM =  512

    # You can experiment with this number to see if it improves training
    STEPS_PER_EPOCH = 300

    # This is how often validation is run. If you are using too much hard drive space
    # on saved models (in the MODEL_DIR), try making this value larger.
    VALIDATION_STEPS = 20
    
    # Matterport originally used resnet101, but I downsized to fit it on my graphics card
    BACKBONE = 'resnet50'

    #RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)
    TRAIN_ROIS_PER_IMAGE = 32
    MAX_GT_INSTANCES = 50 
    POST_NMS_ROIS_INFERENCE = 200 
    POST_NMS_ROIS_TRAINING = 500 
    
config = base_modelConfig()
config.display()

# Configuration of final model
class final_modelConfig(Config):  #
    """Configuration for training on the cigarette butts dataset.
    Derives from the base Config class and overrides values specific
    to the cigarette butts dataset.
    """
    # Give the configuration a recognizable name
    NAME = "plant_phenotyping"

    # Train on 1 GPU and 1 image per GPU. Batch size is 1 (GPUs * images/GPU).
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

    # Number of classes (including background)
    NUM_CLASSES = 1 + 2  # background + 1 (cig_butt)

    # All of our training images are 512x512
    IMAGE_MIN_DIM = 512
    IMAGE_MAX_DIM =  512

    # You can experiment with this number to see if it improves training
    STEPS_PER_EPOCH = 500

    # This is how often validation is run. If you are using too much hard drive space
    # on saved models (in the MODEL_DIR), try making this value larger.
    VALIDATION_STEPS = 5
    
    # Matterport originally used resnet101, but I downsized to fit it on my graphics card
    BACKBONE = 'resnet101'

    # To be honest, I haven't taken the time to figure out what these do
    RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)
    TRAIN_ROIS_PER_IMAGE = 32
    MAX_GT_INSTANCES = 50 
    POST_NMS_ROIS_INFERENCE = 500 
    POST_NMS_ROIS_TRAINING = 1000 
    
final_config = final_modelConfig()
final_config.display()

def get_ax(rows=1, cols=1, size=8):
    """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

class PlantDataset(utils.Dataset):
    """ Generates a COCO-like dataset, i.e. an image dataset annotated in the style of the COCO dataset.
        See http://cocodataset.org/#home for more information.
    """
    def load_data(self, annotation_json, images_dir):
        """ Load the coco-like dataset from json
        Args:
            annotation_json: The path to the coco annotations json file
            images_dir: The directory holding the images referred to by the json file
        """

        # Add all 2 classes
        self.add_class("plant", 1, "leaf")
        self.add_class("plant", 2, "collar")

        coco = COCO(annotation_json)
        # Load all images from the coco.json file
        image_ids = list(coco.imgs.keys())
        # get path to all images
        image_paths = [IMAGE_DIR + coco.imgs[i]['file_name'] for i in image_ids]

        # Add images
        # all images have been resized to 224x224
        for idx, id in enumerate(image_ids):
          self.add_image("plant", image_id=id,
                         path=image_paths[idx],
                         width=512, height=512,
                         annotations=coco.loadAnns(coco.getAnnIds(imgIds=[id],
                                                                  catIds=[1,2],
                                                                  iscrowd=None)))
                
    def load_mask(self, image_id):
        """ Load instance masks for the given image.
        MaskRCNN expects masks in the form of a bitmap [height, width, instances].
        Args:
            image_id: The id of the image to load masks for
        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.
        """
        image_info = self.image_info[image_id]
        annotations = image_info['annotations']
        instance_masks = []
        class_ids = []
        
        for annotation in annotations:
            class_id = annotation['category_id']
            mask = Image.new('1', (image_info['width'], image_info['height']))
            mask_draw = ImageDraw.ImageDraw(mask, '1')
            #for segmentation in annotation['keypoints']:
            
            keypoint = annotation['keypoints']
            keypoint = keypoint[0:4]
            mask_draw.rectangle(keypoint, fill=1, width=2)
            bool_array = np.array(mask) > 0
            instance_masks.append(bool_array)
            class_ids.append(class_id)

        mask = np.dstack(instance_masks)
        class_ids = np.array(class_ids, dtype=np.int32)
        
        return mask, class_ids

"""# Load training and validation datat"""

# Load pre-split training and validation data
print("loading training dataset")
dataset_train = PlantDataset()
dataset_train.load_data(TRAIN_PATH, IMAGE_DIR)
dataset_train.prepare()

print("loading validation dataset")
dataset_val = PlantDataset()
dataset_val.load_data(VAL_PATH, IMAGE_DIR)
dataset_val.prepare()

# visual inspection of data and ground truth
dataset = dataset_train
image_ids = np.random.choice(dataset.image_ids, 4)
for image_id in image_ids:
    image = dataset.load_image(image_id)
    mask, class_ids = dataset.load_mask(image_id)
    visualize.display_top_masks(image, mask, class_ids, dataset.class_names)

"""# Train Mask-RCNN base model"""

# Create model in training mode
model = modellib.MaskRCNN(mode="training", config=base_config,
                          model_dir=MODEL_DIR)

# Which weights to start with?
#init_with = "coco"  # imagenet, coco, or last
init_with = "imagenet"

if init_with == "imagenet":
    model.load_weights(model.get_imagenet_weights(), by_name=True)
elif init_with == "coco":
    # Load weights trained on MS COCO, but skip layers that
    # are different due to the different number of classes
    # See README for instructions to download the COCO weights
    model.load_weights(COCO_MODEL_PATH, by_name=True,
                       exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", 
                                "mrcnn_bbox", "mrcnn_mask"])
elif init_with == "last":
    # Load the last model you trained and continue training
    model.load_weights(model.find_last(), by_name=True)

# Passing layers="heads" freezes all layers except the head
# layers. You can also pass a regular expression to select
# which layers to train by name pattern.
model.train(dataset_train, dataset_val, 
            learning_rate=config.LEARNING_RATE, 
            epochs=2, 
            layers='heads')

# Fine tune all layers
# Passing layers="all" trains all layers. You can also 
# pass a regular expression to select which layers to
# train by name pattern.
model.train(dataset_train, dataset_val, 
            learning_rate=config.LEARNING_RATE / 10,
            epochs=10, 
            layers="all")

"""# plots to check performance"""

hist = model.keras_model.history.history
# plt.plot(hist.history['acc'])
# plt.plot(hist.history['val_acc'])
# plt.title('Model accuracy')
# plt.ylabel('Accuracy')
# plt.xlabel('Epoch')
# plt.legend(['Train', 'Val'], loc='upper left')
# plt.show()

plt.plot(hist['val_loss'])
plt.title('Validation loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
# plt.legend(['Train', 'Val'], loc='upper left')
plt.show()

plt.plot(hist['val_rpn_class_loss'])
plt.title('Validation RPN Class loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.show()

plt.plot(hist['val_rpn_bbox_loss'])
plt.title('Validation RPN Bounding box loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.show()

plt.plot(hist['val_mrcnn_class_loss'])
plt.title('Validation Mask RCNN Class loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.show()

plt.plot(hist['val_mrcnn_bbox_loss'])
plt.title('Validation Mask RCNN Bounding box loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.show()

plt.plot(hist['val_mrcnn_mask_loss'])
plt.title('Validation Mask RCNN mask loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.show()

plt.plot(hist['loss'])
plt.title('Loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.show()

plt.plot(hist['rpn_class_loss'])
plt.title('RPN Class Loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.show()

plt.plot(hist['rpn_bbox_loss'])
plt.title('RPN Bounding box loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.show()

plt.plot(hist['mrcnn_class_loss'])
plt.title('Mask RCNN Class loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.show()

plt.plot(hist['mrcnn_bbox_loss'])
plt.title('Mask RCNN Bounding box loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.show()

plt.plot(hist['mrcnn_mask_loss'])
plt.title('Mask RCNN mask loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.show()

"""# Test model"""

# Code to count the number of leaves and collars detected
def leaf_collar_counts(detect_result):
  d = dict()
  counts = np.unique(detect_result, return_counts=True)
  class_ids = counts[0]
  count = counts[1]
  for idx in range(len(class_ids)):
    d[class_ids[idx]] = count[idx]
  return d

# code to compute mAp evaluation metrics for trained model
def evaluate_model(model, dataset_val, inference_config):
  # Compute VOC-Style mAP @ IoU = 0.5
  image_ids = dataset_val.image_ids
  APs = []
  counts = []
  for image_id in image_ids:
    # Load image and ground truth data
    image, image_meta, gt_class_id, gt_bbox, gt_mask =\
        modellib.load_image_gt(dataset_val, inference_config,
                               image_id, use_mini_mask=False)
    molded_images = np.expand_dims(modellib.mold_image(image, inference_config), 0)

    # Run object detection
    results = model.detect([image], verbose=0)
    r = results[0]

    # Count leaves and collars
    counts.append(leaf_collar_counts(r["class_ids"]))

    # Compute AP
    AP, precisions, recalls, overlaps =\
        utils.compute_ap(gt_bbox, gt_class_id, gt_mask,
                         r["rois"], r["class_ids"], r["scores"], r['masks'])
    APs.append(AP)
    bbox_over = utils.compute_overlaps(gt_bbox, r["rois"])

  return APs, precisions, recalls, counts

class BaseInferenceConfig(base_modelConfig):
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
class FinalInferenceConfig(final_modelConfig):
  GPU_COUNT = 1
  IMAGES_PER_GPU = 1


base_inference_config = BaseInferenceConfig()
final_inference_config = FinalInferenceConfig()

# Recreate the model in inference mode
base_model = modellib.MaskRCNN(mode="inference", 
                          config=base_inference_config,
                          model_dir=MODEL_DIR)
final_model = modellib.MaskRCNN(mode="inference",
                                config=final_inference_config,
                                model_dir=MODEL_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()
base_model_path = os.path.join(ROOT_DIR, "logs/mask_rcnn_plant_phenotyping_0010.h5")
final_model_path = os.path.join(ROOT_DIR, "logs/mask_rcnn_plant_phenotyping_0050.h5")

# Load trained weights of base model
print("Loading weights from ", base_model_path)
base_model.load_weights(base_model_path, by_name=True)

# Load trained weights of final model
print("Loading weights from ", final_model_path)
final_model.load_weights(final_model_path, by_name=True)

# Test base model on a random image
# Display ground truth of a random images
image_id = random.choice(dataset_val.image_ids)
original_image, image_meta, gt_class_id, gt_bbox, gt_mask =\
    modellib.load_image_gt(dataset_val, base_inference_config, 
                           image_id, use_mini_mask=False)

log("original_image", original_image)
log("image_meta", image_meta)
log("gt_class_id", gt_class_id)
log("gt_bbox", gt_bbox)
log("gt_mask", gt_mask)

visualize.display_instances(original_image, gt_bbox, gt_mask, gt_class_id, 
                            dataset_train.class_names, figsize=(8, 8))
#print(gt_bbox)

# base model detection result of 1 random image for visualization
base_results = base_model.detect([original_image], verbose=1)
r_base = base_results[0]

count = leaf_collar_counts(r_base["class_ids"])
print("number of leaves: ", count[1])
print("number of collars: ", count[2] )
visualize.display_instances(original_image, r_base['rois'], r_base['masks'], r_base['class_ids'], 
                            dataset_val.class_names, r_base['scores'], ax=get_ax())

final_results = final_model.detect([original_image], verbose=1)
r_final = final_results[0]
count = leaf_collar_counts(r_final["class_ids"])
print("number of leaves: ", count[1])
print("number of collars: ", count[2] )

visualize.display_instances(original_image, r_final['rois'], r_final['masks'], r_final['class_ids'], 
                            dataset_val.class_names, r_final['scores'], ax=get_ax())

# Check the performance of base and final model on the entire validation dataset
AP_base, precisions_base, recalls_base, counts_base = evaluate_model(base_model, dataset_val, base_inference_config)
AP_final, precisions_final, recalls_final, counts_final = evaluate_model(final_model, dataset_val, final_inference_config)

print(np.mean(AP_base))
print(np.mean(AP_final))

print(counts_base)

"""# Plot test performance"""

visualize.plot_precision_recall(AP, precisions, recalls)
visualize.plot_overlaps(gt_class_id, r['class_ids'], r['scores'],
                        overlaps, dataset.class_names)