<reponame>SushmaDG/MaskRCNN
import numpy as np
import cv2
from mrcnn.config import Config
from mrcnn.model as MaskRCNN
from shapes_loader import Shapes

MODEL_DIR = '/gluster/home/sdevaramani/Thesis'
COCO_MODEL_PATH = '/gluster/home/sdevaramani/Thesis/weights/mask_rcnn_coco.h5'


class ShapesConfig(Config):
    """Configuration for training on the toy shapes dataset.
    """
    NAME = 'shapes'
    GPU_COUNT = 1
    IMAGES_PER_GPU = 8
    NUM_CLASSES = 1 + 3  # background + 3 shapes
    IMAGE_MIN_DIM = 128
    IMAGE_MAX_DIM = 128
    RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)
    TRAIN_ROIS_PER_IMAGE = 32
    STEPS_PER_EPOCH = 100
    VALIDATION_STEPS = 5

config = ShapesConfig()

# Training dataset
dataset_train = Shapes()
dataset_train.load_shapes(500, config.IMAGE_SHAPE[0], config.IMAGE_SHAPE[1])
dataset_train.prepare()

# Validation dataset
dataset_val = Shapes()
dataset_val.load_shapes(50, config.IMAGE_SHAPE[0], config.IMAGE_SHAPE[1])
dataset_val.prepare()

# Create model in training mode
model = MaskRCNN(mode='training', config=config,
                 model_dir=MODEL_DIR)

init_with = 'coco'  # imagenet, coco, or last

if init_with == 'imagenet':
    model.load_weights(model.get_imagenet_weights(), by_name=True)
elif init_with == 'coco':
    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)

model.train(dataset_train, dataset_val, 
            learning_rate=config.LEARNING_RATE, 
            epochs=1, 
            layers='heads')

model.train(dataset_train, dataset_val, 
            learning_rate=config.LEARNING_RATE / 10,
            epochs=2, 
            layers='all')
