import os

import sys

import random

import math

import re

import time

import numpy as np

import cv2

import matplotlib

import matplotlib.pyplot as plt

from PIL import Image

# Root directory of the project

ROOT_DIR = os.path.abspath("../../")

# Import Mask RCNN

sys.path.append(ROOT_DIR)  # To find local version of the library

from mrcnn.config import Config

from mrcnn import utils

import mrcnn.model as modellib

from mrcnn import visualize

from mrcnn.model import log

#%matplotlib inline 

# Directory to save logs and trained model

MODEL_DIR = os.path.join(ROOT_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)

iter_num=0

class ShapesConfig(Config):

    """Configuration for training on the toy shapes dataset.

    Derives from the base Config class and overrides values specific

    to the toy shapes dataset.

    """

    # Give the configuration a recognizable name

    NAME = "shapes"

    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each

    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).

    GPU_COUNT = 2

    IMAGES_PER_GPU = 1 #这里我用了两个GPU

    # Number of classes (including background)

    NUM_CLASSES = 1 + 1  # background + 1 shapes

    # Use small images for faster training. Set the limits of the small side

    # the large side, and that determines the image shape.

    IMAGE_MIN_DIM = 1080

    IMAGE_MAX_DIM = 1920

    # Use smaller anchors because our image and objects are small

    RPN_ANCHOR_SCALES = (8*6, 16*6, 32*6, 64*6, 128*6)  # anchor side in pixels

    # Reduce training ROIs per image because the images are small and have

    # few objects. Aim to allow ROI sampling to pick 33% positive ROIs.

    TRAIN_ROIS_PER_IMAGE = 32

    # Use a small epoch since the data is simple

    STEPS_PER_EPOCH = 100

    # use small validation steps since the epoch is small

    VALIDATION_STEPS = 5

config = ShapesConfig()

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 DrugDataset(utils.Dataset):

    #得到该图中有多少个实例（物体）

    def get_obj_index(self, image):

        n = np.max(image)

        return n

    #解析labelme中得到的yaml文件，从而得到mask每一层对应的实例标签

    def from_yaml_get_class(self,image_id):

        info=self.image_info[image_id]

        with open(info['yaml_path']) as f:

            temp=yaml.load(f.read())

            labels=temp['label_names']

            del labels[0]

        return labels

    #重新写draw_mask

    def draw_mask(self, num_obj, mask, image):

        info = self.image_info[image_id]

        for index in range(num_obj):

            for i in range(info['width']):

                for j in range(info['height']):

                    at_pixel = image.getpixel((i, j))

                    if at_pixel == index + 1:

                        mask[j, i, index] =1

        return mask

    #重新写load_shapes，里面包含自己的自己的类别（我的是box、column、package、fruit四类）

    #并在self.image_info信息中添加了path、mask_path 、yaml_path

    def load_shapes(self, count, height, width, img_floder, mask_floder, imglist,dataset_root_path):

        """Generate the requested number of synthetic images.

        count: number of images to generate.

        height, width: the size of the generated images.

        """

        # Add classes

        self.add_class("shapes", 1, "box")

        for i in range(count):

            filestr = imglist[i].split(".")[0]

            filestr = filestr.split("_")[0]

            mask_path = mask_floder + "/" + filestr + ".png"

            yaml_path=dataset_root_path+filestr+"rgb_"+"_json/info.yaml"

            self.add_image("shapes", image_id=i, path=img_floder + "/"+imglist[i],

                           width=width, height=height, mask_path=mask_path,yaml_path=yaml_path)

    #重写load_mask

    def load_mask(self, image_id):

        """Generate instance masks for shapes of the given image ID.

        """

        global iter_num

        info = self.image_info[image_id]

        count = 1  # number of object

        img = Image.open(info['mask_path'])

        num_obj = self.get_obj_index(img)

        mask = np.zeros([info['height'], info['width'], num_obj], dtype=np.uint8)

        mask = self.draw_mask(num_obj, mask, img)

        occlusion = np.logical_not(mask[:, :, -1]).astype(np.uint8)

        for i in range(count - 2, -1, -1):

            mask[:, :, i] = mask[:, :, i] * occlusion

            occlusion = np.logical_and(occlusion, np.logical_not(mask[:, :, i]))

        labels=[]

        labels=self.from_yaml_get_class(image_id)

        labels_form=[]

        for i in range(len(labels)):

            if labels[i].find("box")!=-1:

                #print "box"

                labels_form.append("box")

            #elif labels[i].find("column")!=-1:

                #print "column"

             #   labels_form.append("column")

            #elif labels[i].find("package")!=-1:

                #print "package"

             #   labels_form.append("package")

            #elif labels[i].find("fruit")!=-1:

                #print "fruit"

             #   labels_form.append("fruit")

        class_ids = np.array([self.class_names.index(s) for s in labels_form])

        return mask, class_ids.astype(np.int32)

#基础设置

dataset_root_path="/mnt/disk2/zhouqiang/Mask_RCNN/data/train_01_01/"

img_floder = dataset_root_path+"rgb"

mask_floder = dataset_root_path+"mask"

#yaml_floder = dataset_root_path

imglist = os.listdir(img_floder)

count = len(imglist)

width = 1920

height = 1080

#train与val数据集准备

dataset_train = DrugDataset()

dataset_train.load_shapes(count, 1080, 1920, img_floder, mask_floder, imglist,dataset_root_path)

dataset_train.prepare()

dataset_val = DrugDataset()

dataset_val.load_shapes(count, 1080, 1920, img_floder, mask_floder, imglist,dataset_root_path)

dataset_val.prepare()

# Create model in training mode

model = modellib.MaskRCNN(mode="training", config=config,

                          model_dir=MODEL_DIR)

# Which weights to start with?

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

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

# 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=50,

            layers="all")