"""
Mask R-CNN
The main Mask R-CNN model implementation.

Copyright (c) 2017 Matterport, Inc.
Licensed under the MIT License (see LICENSE for details)
Written by Waleed Abdulla
"""

import os
import random
import datetime
import re
import math
import logging
import sys
from collections import OrderedDict
import multiprocessing
import numpy as np
import tensorflow as tf

from mrcnn import utils
from mrcnn.net_util import *
from skimage.transform import resize
import time


############################################################
#  Utility Functions
############################################################
def log(text, array=None):
    """Prints a text message. And, optionally, if a Numpy array is provided it
    prints it's shape, min, and max values.
    """
    if array is not None:
        text = text.ljust(25)
        text += ("shape: {:20}  ".format(str(array.shape)))
        if array.size:
            text += ("min: {:10.5f}  max: {:10.5f}".format(array.min(), array.max()))
        else:
            text += ("min: {:10}  max: {:10}".format("", ""))
        text += "  {}".format(array.dtype)
    print(text)


def compute_backbone_shapes(config, image_shape):
    """Computes the width and height of each stage of the backbone network.

    Returns:
        [N, (height, width)]. Where N is the number of stages
    """
    if callable(config.BACKBONE):  # 现在的使用BACKBONE是一个list，不是callable的，不过也说可以自己实现一个
        return config.COMPUTE_BACKBONE_SHAPE(image_shape)

    # Currently supports ResNet only
    assert config.BACKBONE in ["resnet50", "resnet101"]
    return np.array(
        [[int(math.ceil(image_shape[0] / stride)),
            int(math.ceil(image_shape[1] / stride))]
            for stride in config.BACKBONE_STRIDES])
    # BACKBONE_STRIDES [4, 8, 16, 32, 64]
    # 返回值应该是[[256,256],
    #            [128,128],
    #            [64, 64],
    #            [32, 32],
    #            [16, 16]]


############################################################
#  Data Generator
############################################################
def load_image_gt(dataset, config, image_id, augment=False, augmentation=None,
                  use_mini_mask=False):
    """Load and return ground truth data for an image (image, mask, bounding boxes).

    augment: (deprecated. Use augmentation instead). If true, apply random
        image augmentation. Currently, only horizontal flipping is offered.
    augmentation: Optional. An imgaug (https://github.com/aleju/imgaug) augmentation.
        For example, passing imgaug.augmenters.Fliplr(0.5) flips images
        right/left 50% of the time.
    use_mini_mask: If False, returns full-size masks that are the same height
        and width as the original image. These can be big, for example
        1024x1024x100 (for 100 instances). Mini masks are smaller, typically,
        224x224 and are generated by extracting the bounding box of the
        object and resizing it to MINI_MASK_SHAPE.

    Returns:
    image: [height, width, 3]
    shape: the original shape of the image before resizing and cropping.
    class_ids: [instance_count] Integer class IDs
    bbox: [instance_count, (y1, x1, y2, x2)]
    mask: [height, width, instance_count]. The height and width are those
        of the image unless use_mini_mask is True, in which case they are
        defined in MINI_MASK_SHAPE.
    """
    # Load image and mask
    image = dataset.load_image(image_id)
    mask, class_ids = dataset.load_mask(image_id)
    original_shape = image.shape
    image, window, scale, padding, crop = utils.resize_image(
        image,
        min_dim=config.IMAGE_MIN_DIM,
        min_scale=config.IMAGE_MIN_SCALE,
        max_dim=config.IMAGE_MAX_DIM,
        mode=config.IMAGE_RESIZE_MODE)
    mask = utils.resize_mask(mask, scale, padding, crop)

    # Random horizontal flips.
    # TODO: will be removed in a future update in favor of augmentation
    if augment:
        logging.warning("'augment' is deprecated. Use 'augmentation' instead.")
        if random.randint(0, 1):
            image = np.fliplr(image)
            mask = np.fliplr(mask)

    # Augmentation
    # This requires the imgaug lib (https://github.com/aleju/imgaug)
    if augmentation:
        import imgaug

        # Augmenters that are safe to apply to masks
        # Some, such as Affine, have settings that make them unsafe, so always
        # test your augmentation on masks
        MASK_AUGMENTERS = ["Sequential", "SomeOf", "OneOf", "Sometimes",
                           "Fliplr", "Flipud", "CropAndPad",
                           "Affine", "PiecewiseAffine"]

        def hook(images, augmenter, parents, default):
            """Determines which augmenters to apply to masks."""
            return augmenter.__class__.__name__ in MASK_AUGMENTERS

        # Store shapes before augmentation to compare
        image_shape = image.shape
        mask_shape = mask.shape
        # Make augmenters deterministic to apply similarly to images and masks
        det = augmentation.to_deterministic()
        image = det.augment_image(image)
        # Change mask to np.uint8 because imgaug doesn't support np.bool
        mask = det.augment_image(mask.astype(np.uint8),
                                 hooks=imgaug.HooksImages(activator=hook))
        # Verify that shapes didn't change
        assert image.shape == image_shape, "Augmentation shouldn't change image size"
        assert mask.shape == mask_shape, "Augmentation shouldn't change mask size"
        # Change mask back to bool
        mask = mask.astype(np.bool)

    # Note that some boxes might be all zeros if the corresponding mask got cropped out.
    # and here is to filter them out
    _idx = np.sum(mask, axis=(0, 1)) > 0
    mask = mask[:, :, _idx]
    class_ids = class_ids[_idx]
    # Bounding boxes. Note that some boxes might be all zeros
    # if the corresponding mask got cropped out.
    # bbox: [num_instances, (y1, x1, y2, x2)]
    bbox = utils.extract_bboxes(mask)

    # Active classes
    # Different datasets have different classes, so track the
    # classes supported in the dataset of this image.
    active_class_ids = np.zeros([dataset.num_classes], dtype=np.int32)
    source_class_ids = dataset.source_class_ids[dataset.image_info[image_id]["source"]]
    active_class_ids[source_class_ids] = 1

    # Resize masks to smaller size to reduce memory usage
    if use_mini_mask:
        mask = utils.minimize_mask(bbox, mask, config.MINI_MASK_SHAPE)

    # Image meta data
    image_meta = compose_image_meta(image_id, original_shape, image.shape,
                                    window, scale, active_class_ids)

    return image, image_meta, class_ids, bbox, mask


def build_detection_targets(rpn_rois, gt_class_ids, gt_boxes, gt_masks, config):
    """Generate targets for training Stage 2 classifier and mask heads.
    This is not used in normal training. It's useful for debugging or to train
    the Mask RCNN heads without using the RPN head.

    Inputs:
    rpn_rois: [N, (y1, x1, y2, x2)] proposal boxes.
    gt_class_ids: [instance count] Integer class IDs
    gt_boxes: [instance count, (y1, x1, y2, x2)]
    gt_masks: [height, width, instance count] Ground truth masks. Can be full
              size or mini-masks.

    Returns:
    rois: [TRAIN_ROIS_PER_IMAGE, (y1, x1, y2, x2)]
    class_ids: [TRAIN_ROIS_PER_IMAGE]. Integer class IDs.
    bboxes: [TRAIN_ROIS_PER_IMAGE, NUM_CLASSES, (y, x, log(h), log(w))]. Class-specific
            bbox refinements.
    masks: [TRAIN_ROIS_PER_IMAGE, height, width, NUM_CLASSES). Class specific masks cropped
           to bbox boundaries and resized to neural network output size.
    """
    assert rpn_rois.shape[0] > 0
    assert gt_class_ids.dtype == np.int32, "Expected int but got {}".format(
        gt_class_ids.dtype)
    assert gt_boxes.dtype == np.int32, "Expected int but got {}".format(
        gt_boxes.dtype)
    assert gt_masks.dtype == np.bool_, "Expected bool but got {}".format(
        gt_masks.dtype)

    # It's common to add GT Boxes to ROIs but we don't do that here because
    # according to XinLei Chen's paper, it doesn't help.

    # Trim empty padding in gt_boxes and gt_masks parts
    instance_ids = np.where(gt_class_ids > 0)[0]
    assert instance_ids.shape[0] > 0, "Image must contain instances."
    gt_class_ids = gt_class_ids[instance_ids]
    gt_boxes = gt_boxes[instance_ids]
    gt_masks = gt_masks[:, :, instance_ids]

    # Compute areas of ROIs and ground truth boxes.
    rpn_roi_area = (rpn_rois[:, 2] - rpn_rois[:, 0]) * \
        (rpn_rois[:, 3] - rpn_rois[:, 1])
    gt_box_area = (gt_boxes[:, 2] - gt_boxes[:, 0]) * \
        (gt_boxes[:, 3] - gt_boxes[:, 1])

    # Compute overlaps [rpn_rois, gt_boxes]
    overlaps = np.zeros((rpn_rois.shape[0], gt_boxes.shape[0]))
    for i in range(overlaps.shape[1]):
        gt = gt_boxes[i]
        overlaps[:, i] = utils.compute_iou(
            gt, rpn_rois, gt_box_area[i], rpn_roi_area)

    # Assign ROIs to GT boxes
    rpn_roi_iou_argmax = np.argmax(overlaps, axis=1)
    rpn_roi_iou_max = overlaps[np.arange(
        overlaps.shape[0]), rpn_roi_iou_argmax]
    # GT box assigned to each ROI
    rpn_roi_gt_boxes = gt_boxes[rpn_roi_iou_argmax]
    rpn_roi_gt_class_ids = gt_class_ids[rpn_roi_iou_argmax]

    # Positive ROIs are those with >= 0.5 IoU with a GT box.
    fg_ids = np.where(rpn_roi_iou_max > 0.5)[0]

    # Negative ROIs are those with max IoU 0.1-0.5 (hard example mining)
    # TODO: To hard example mine or not to hard example mine, that's the question
    # bg_ids = np.where((rpn_roi_iou_max >= 0.1) & (rpn_roi_iou_max < 0.5))[0]
    bg_ids = np.where(rpn_roi_iou_max < 0.5)[0]

    # Subsample ROIs. Aim for 33% foreground.
    # FG
    fg_roi_count = int(config.TRAIN_ROIS_PER_IMAGE * config.ROI_POSITIVE_RATIO)
    if fg_ids.shape[0] > fg_roi_count:
        keep_fg_ids = np.random.choice(fg_ids, fg_roi_count, replace=False)
    else:
        keep_fg_ids = fg_ids
    # BG
    remaining = config.TRAIN_ROIS_PER_IMAGE - keep_fg_ids.shape[0]
    if bg_ids.shape[0] > remaining:
        keep_bg_ids = np.random.choice(bg_ids, remaining, replace=False)
    else:
        keep_bg_ids = bg_ids
    # Combine indices of ROIs to keep
    keep = np.concatenate([keep_fg_ids, keep_bg_ids])
    # Need more?
    remaining = config.TRAIN_ROIS_PER_IMAGE - keep.shape[0]
    if remaining > 0:
        # Looks like we don't have enough samples to maintain the desired
        # balance. Reduce requirements and fill in the rest. This is
        # likely different from the Mask RCNN paper.

        # There is a small chance we have neither fg nor bg samples.
        if keep.shape[0] == 0:
            # Pick bg regions with easier IoU threshold
            bg_ids = np.where(rpn_roi_iou_max < 0.5)[0]
            assert bg_ids.shape[0] >= remaining
            keep_bg_ids = np.random.choice(bg_ids, remaining, replace=False)
            assert keep_bg_ids.shape[0] == remaining
            keep = np.concatenate([keep, keep_bg_ids])
        else:
            # Fill the rest with repeated bg rois.
            keep_extra_ids = np.random.choice(
                keep_bg_ids, remaining, replace=True)
            keep = np.concatenate([keep, keep_extra_ids])
    assert keep.shape[0] == config.TRAIN_ROIS_PER_IMAGE, \
        "keep doesn't match ROI batch size {}, {}".format(
            keep.shape[0], config.TRAIN_ROIS_PER_IMAGE)

    # Reset the gt boxes assigned to BG ROIs.
    rpn_roi_gt_boxes[keep_bg_ids, :] = 0
    rpn_roi_gt_class_ids[keep_bg_ids] = 0

    # For each kept ROI, assign a class_id, and for FG ROIs also add bbox refinement.
    rois = rpn_rois[keep]
    roi_gt_boxes = rpn_roi_gt_boxes[keep]
    roi_gt_class_ids = rpn_roi_gt_class_ids[keep]
    roi_gt_assignment = rpn_roi_iou_argmax[keep]

    # Class-aware bbox deltas. [y, x, log(h), log(w)]
    bboxes = np.zeros((config.TRAIN_ROIS_PER_IMAGE,
                       config.NUM_CLASSES, 4), dtype=np.float32)
    pos_ids = np.where(roi_gt_class_ids > 0)[0]
    bboxes[pos_ids, roi_gt_class_ids[pos_ids]] = utils.box_refinement(
        rois[pos_ids], roi_gt_boxes[pos_ids, :4])
    # Normalize bbox refinements
    bboxes /= config.BBOX_STD_DEV

    # Generate class-specific target masks
    masks = np.zeros((config.TRAIN_ROIS_PER_IMAGE, config.MASK_SHAPE[0], config.MASK_SHAPE[1], config.NUM_CLASSES),
                     dtype=np.float32)
    for i in pos_ids:
        class_id = roi_gt_class_ids[i]
        assert class_id > 0, "class id must be greater than 0"
        gt_id = roi_gt_assignment[i]
        class_mask = gt_masks[:, :, gt_id]

        if config.USE_MINI_MASK:
            # Create a mask placeholder, the size of the image
            placeholder = np.zeros(config.IMAGE_SHAPE[:2], dtype=bool)
            # GT box
            gt_y1, gt_x1, gt_y2, gt_x2 = gt_boxes[gt_id]
            gt_w = gt_x2 - gt_x1
            gt_h = gt_y2 - gt_y1
            # Resize mini mask to size of GT box
            placeholder[gt_y1:gt_y2, gt_x1:gt_x2] = \
                np.round(utils.resize(class_mask, (gt_h, gt_w))).astype(bool)
            # Place the mini batch in the placeholder
            class_mask = placeholder

        # Pick part of the mask and resize it
        y1, x1, y2, x2 = rois[i].astype(np.int32)
        m = class_mask[y1:y2, x1:x2]
        mask = utils.resize(m, config.MASK_SHAPE)
        masks[i, :, :, class_id] = mask

    return rois, roi_gt_class_ids, bboxes, masks


def build_rpn_targets(image_shape, anchors, gt_class_ids, gt_boxes, config):
    """Given the anchors and GT boxes, compute overlaps and identify positive
    anchors and deltas to refine them to match their corresponding GT boxes.

    anchors: [num_anchors, (y1, x1, y2, x2)]
    gt_class_ids: [num_gt_boxes] Integer class IDs.
    gt_boxes: [num_gt_boxes, (y1, x1, y2, x2)]

    Returns:
    rpn_match: [N] (int32) matches between anchors and GT boxes.
               1 = positive anchor, -1 = negative anchor, 0 = neutral
    rpn_bbox: [N, (dy, dx, log(dh), log(dw))] Anchor bbox deltas.
    """
    # RPN Match: 1 = positive anchor, -1 = negative anchor, 0 = neutral
    rpn_match = np.zeros([anchors.shape[0]], dtype=np.int32)
    # RPN bounding boxes: [max anchors per image, (dy, dx, log(dh), log(dw))]
    rpn_bbox = np.zeros((config.RPN_TRAIN_ANCHORS_PER_IMAGE, 4))

    # Handle COCO crowds
    # A crowd box in COCO is a bounding box around several instances. Exclude
    # them from training. A crowd box is given a negative class ID.
    crowd_ix = np.where(gt_class_ids < 0)[0]
    if crowd_ix.shape[0] > 0:
        # Filter out crowds from ground truth class IDs and boxes
        non_crowd_ix = np.where(gt_class_ids > 0)[0]
        crowd_boxes = gt_boxes[crowd_ix]
        gt_class_ids = gt_class_ids[non_crowd_ix]
        gt_boxes = gt_boxes[non_crowd_ix]
        # Compute overlaps with crowd boxes [anchors, crowds]
        crowd_overlaps = utils.compute_overlaps(anchors, crowd_boxes)
        crowd_iou_max = np.amax(crowd_overlaps, axis=1)
        no_crowd_bool = (crowd_iou_max < 0.001)
    else:
        # All anchors don't intersect a crowd
        no_crowd_bool = np.ones([anchors.shape[0]], dtype=bool)

    # Compute overlaps [num_anchors, num_gt_boxes]
    overlaps = utils.compute_overlaps(anchors, gt_boxes)

    # Match anchors to GT Boxes
    # If an anchor overlaps a GT box with IoU >= 0.7 then it's positive.
    # If an anchor overlaps a GT box with IoU < 0.3 then it's negative.
    # Neutral anchors are those that don't match the conditions above,
    # and they don't influence the loss function.
    # However, don't keep any GT box unmatched (rare, but happens). Instead,
    # match it to the closest anchor (even if its max IoU is < 0.3).
    #
    # 1. Set negative anchors first. They get overwritten below if a GT box is
    # matched to them. Skip boxes in crowd areas.
    anchor_iou_argmax = np.argmax(overlaps, axis=1)
    anchor_iou_max = overlaps[np.arange(overlaps.shape[0]), anchor_iou_argmax]
    rpn_match[(anchor_iou_max < 0.3) & no_crowd_bool] = -1
    # 2. Set an anchor for each GT box (regardless of IoU value).
    # If multiple anchors have the same IoU match all of them
    gt_iou_argmax = np.argwhere(overlaps == np.max(overlaps, axis=0))[:, 0]
    rpn_match[gt_iou_argmax] = 1
    # 3. Set anchors with high overlap as positive.
    rpn_match[anchor_iou_max >= 0.7] = 1

    # Subsample to balance positive and negative anchors
    # Don't let positives be more than half the anchors
    ids = np.where(rpn_match == 1)[0]
    extra = len(ids) - (config.RPN_TRAIN_ANCHORS_PER_IMAGE // 2)
    if extra > 0:
        # Reset the extra ones to neutral
        ids = np.random.choice(ids, extra, replace=False)  # 从ids中随机选出extra个数
        rpn_match[ids] = 0
    # Same for negative proposals
    ids = np.where(rpn_match == -1)[0]
    extra = len(ids) - (config.RPN_TRAIN_ANCHORS_PER_IMAGE -
                        np.sum(rpn_match == 1))
    if extra > 0:
        # Rest the extra ones to neutral
        ids = np.random.choice(ids, extra, replace=False)
        rpn_match[ids] = 0

    # For positive anchors, compute shift and scale needed to transform them
    # to match the corresponding GT boxes.
    ids = np.where(rpn_match == 1)[0]
    ix = 0  # index into rpn_bbox
    # TODO: use box_refinement() rather than duplicating the code here
    for i, a in zip(ids, anchors[ids]):
        # Closest gt box (it might have IoU < 0.7)
        gt = gt_boxes[anchor_iou_argmax[i]]

        # Convert coordinates to center plus width/height.
        # GT Box
        gt_h = gt[2] - gt[0]
        gt_w = gt[3] - gt[1]
        gt_center_y = gt[0] + 0.5 * gt_h
        gt_center_x = gt[1] + 0.5 * gt_w
        # Anchor
        a_h = a[2] - a[0]
        a_w = a[3] - a[1]
        a_center_y = a[0] + 0.5 * a_h
        a_center_x = a[1] + 0.5 * a_w

        # Compute the bbox refinement that the RPN should predict.
        rpn_bbox[ix] = [
            (gt_center_y - a_center_y) / a_h,
            (gt_center_x - a_center_x) / a_w,
            np.log(gt_h / a_h),
            np.log(gt_w / a_w),
        ]
        # Normalize
        rpn_bbox[ix] /= config.RPN_BBOX_STD_DEV
        ix += 1

    return rpn_match, rpn_bbox


def generate_random_rois(image_shape, count, gt_class_ids, gt_boxes):
    """Generates ROI proposals similar to what a region proposal network
    would generate.

    image_shape: [Height, Width, Depth]
    count: Number of ROIs to generate
    gt_class_ids: [N] Integer ground truth class IDs
    gt_boxes: [N, (y1, x1, y2, x2)] Ground truth boxes in pixels.

    Returns: [count, (y1, x1, y2, x2)] ROI boxes in pixels.
    """
    # placeholder
    rois = np.zeros((count, 4), dtype=np.int32)

    # Generate random ROIs around GT boxes (90% of count)
    rois_per_box = int(0.9 * count / gt_boxes.shape[0])
    for i in range(gt_boxes.shape[0]):
        gt_y1, gt_x1, gt_y2, gt_x2 = gt_boxes[i]
        h = gt_y2 - gt_y1
        w = gt_x2 - gt_x1
        # random boundaries
        r_y1 = max(gt_y1 - h, 0)
        r_y2 = min(gt_y2 + h, image_shape[0])
        r_x1 = max(gt_x1 - w, 0)
        r_x2 = min(gt_x2 + w, image_shape[1])

        # To avoid generating boxes with zero area, we generate double what
        # we need and filter out the extra. If we get fewer valid boxes
        # than we need, we loop and try again.
        while True:
            y1y2 = np.random.randint(r_y1, r_y2, (rois_per_box * 2, 2))
            x1x2 = np.random.randint(r_x1, r_x2, (rois_per_box * 2, 2))
            # Filter out zero area boxes
            threshold = 1
            y1y2 = y1y2[np.abs(y1y2[:, 0] - y1y2[:, 1]) >=
                        threshold][:rois_per_box]
            x1x2 = x1x2[np.abs(x1x2[:, 0] - x1x2[:, 1]) >=
                        threshold][:rois_per_box]
            if y1y2.shape[0] == rois_per_box and x1x2.shape[0] == rois_per_box:
                # 这句话感觉有点多余啊，上面都[:rois_per_box]，肯定相等啊
                break

        # Sort on axis 1 to ensure x1 <= x2 and y1 <= y2 and then reshape
        # into x1, y1, x2, y2 order
        x1, x2 = np.split(np.sort(x1x2, axis=1), 2, axis=1)
        y1, y2 = np.split(np.sort(y1y2, axis=1), 2, axis=1)
        box_rois = np.hstack([y1, x1, y2, x2])
        rois[rois_per_box * i:rois_per_box * (i + 1)] = box_rois

    # Generate random ROIs anywhere in the image (10% of count)
    remaining_count = count - (rois_per_box * gt_boxes.shape[0])
    # To avoid generating boxes with zero area, we generate double what
    # we need and filter out the extra. If we get fewer valid boxes
    # than we need, we loop and try again.
    while True:
        y1y2 = np.random.randint(0, image_shape[0], (remaining_count * 2, 2))
        x1x2 = np.random.randint(0, image_shape[1], (remaining_count * 2, 2))
        # Filter out zero area boxes
        threshold = 1
        y1y2 = y1y2[np.abs(y1y2[:, 0] - y1y2[:, 1]) >=
                    threshold][:remaining_count]
        x1x2 = x1x2[np.abs(x1x2[:, 0] - x1x2[:, 1]) >=
                    threshold][:remaining_count]
        if y1y2.shape[0] == remaining_count and x1x2.shape[0] == remaining_count:
            break

    # Sort on axis 1 to ensure x1 <= x2 and y1 <= y2 and then reshape
    # into x1, y1, x2, y2 order
    x1, x2 = np.split(np.sort(x1x2, axis=1), 2, axis=1)
    y1, y2 = np.split(np.sort(y1y2, axis=1), 2, axis=1)
    global_rois = np.hstack([y1, x1, y2, x2])
    rois[-remaining_count:] = global_rois
    return rois


def data_generator(dataset, config, shuffle=True, augment=False, augmentation=None,
                   random_rois=0, batch_size=1, detection_targets=False,
                   no_augmentation_sources=None):
    """A generator that returns images and corresponding target class ids,
    bounding box deltas, and masks.

    dataset: The Dataset object to pick data from
    config: The model config object
    shuffle: If True, shuffles the samples before every epoch
    augment: (deprecated. Use augmentation instead). If true, apply random
        image augmentation. Currently, only horizontal flipping is offered.
    augmentation: Optional. An imgaug (https://github.com/aleju/imgaug) augmentation.
        For example, passing imgaug.augmenters.Fliplr(0.5) flips images
        right/left 50% of the time.
    random_rois: If > 0 then generate proposals to be used to train the
                 network classifier and mask heads. Useful if training
                 the Mask RCNN part without the RPN.
    batch_size: How many images to return in each call
    detection_targets: If True, generate detection targets (class IDs, bbox
        deltas, and masks). Typically for debugging or visualizations because
        in trainig detection targets are generated by DetectionTargetLayer.
    no_augmentation_sources: Optional. List of sources to exclude for
        augmentation. A source is string that identifies a dataset and is
        defined in the Dataset class.

    Returns a Python generator. Upon calling next() on it, the
    generator returns two lists, inputs and outputs. The contents
    of the lists differs depending on the received arguments:
    inputs list:
    - images: [batch, H, W, C]
    - image_meta: [batch, (meta data)] Image details. See compose_image_meta()
    - rpn_match: [batch, N] Integer (1=positive anchor, -1=negative, 0=neutral)
    - rpn_bbox: [batch, N, (dy, dx, log(dh), log(dw))] Anchor bbox deltas.
    - gt_class_ids: [batch, MAX_GT_INSTANCES] Integer class IDs
    - gt_boxes: [batch, MAX_GT_INSTANCES, (y1, x1, y2, x2)]
    - gt_masks: [batch, height, width, MAX_GT_INSTANCES]. The height and width
                are those of the image unless use_mini_mask is True, in which
                case they are defined in MINI_MASK_SHAPE.

    outputs list: Usually empty in regular training. But if detection_targets
        is True then the outputs list contains target class_ids, bbox deltas,
        and masks.
    """
    b = 0  # batch item index
    image_index = -1
    image_ids = np.copy(dataset.image_ids)
    error_count = 0
    no_augmentation_sources = no_augmentation_sources or []

    # Anchors
    # [anchor_count, (y1, x1, y2, x2)]
    backbone_shapes = compute_backbone_shapes(config, config.IMAGE_SHAPE)
    anchors = utils.generate_pyramid_anchors(config.RPN_ANCHOR_SCALES,
                                             config.RPN_ANCHOR_RATIOS,
                                             backbone_shapes,
                                             config.BACKBONE_STRIDES,
                                             config.RPN_ANCHOR_STRIDE)

    # Keras requires a generator to run indefinitely.
    while True:
        try:
            # Increment index to pick next image. Shuffle if at the start of an epoch.
            image_index = (image_index + 1) % len(image_ids)
            if shuffle and image_index == 0:
                np.random.shuffle(image_ids)

            # Get GT bounding boxes and masks for image.
            image_id = image_ids[image_index]

            # If the image source is not to be augmented pass None as augmentation
            if dataset.image_info[image_id]['source'] in no_augmentation_sources:
                image, image_meta, gt_class_ids, gt_boxes, gt_masks = \
                    load_image_gt(dataset, config, image_id, augment=augment,
                                  augmentation=None,
                                  use_mini_mask=config.USE_MINI_MASK)
            else:
                image, image_meta, gt_class_ids, gt_boxes, gt_masks = \
                    load_image_gt(dataset, config, image_id, augment=augment,
                                  augmentation=augmentation,
                                  use_mini_mask=config.USE_MINI_MASK)

            # Skip images that have no instances. This can happen in cases
            # where we train on a subset of classes and the image doesn't
            # have any of the classes we care about.
            if not np.any(gt_class_ids > 0):
                continue

            # RPN Targets
            rpn_match, rpn_bbox = build_rpn_targets(image.shape, anchors,
                                                    gt_class_ids, gt_boxes, config)

            # Mask R-CNN Targets
            if random_rois:
                rpn_rois = generate_random_rois(
                    image.shape, random_rois, gt_class_ids, gt_boxes)
                if detection_targets:
                    rois, mrcnn_class_ids, mrcnn_bbox, mrcnn_mask =\
                        build_detection_targets(
                            rpn_rois, gt_class_ids, gt_boxes, gt_masks, config)

            # Init batch arrays
            if b == 0:
                batch_image_meta = np.zeros(
                    (batch_size,) + image_meta.shape, dtype=image_meta.dtype)
                batch_rpn_match = np.zeros(
                    [batch_size, anchors.shape[0], 1], dtype=rpn_match.dtype)
                batch_rpn_bbox = np.zeros(
                    [batch_size, config.RPN_TRAIN_ANCHORS_PER_IMAGE, 4], dtype=rpn_bbox.dtype)
                batch_images = np.zeros(
                    (batch_size,) + image.shape, dtype=np.float32)
                batch_gt_class_ids = np.zeros(
                    (batch_size, config.MAX_GT_INSTANCES), dtype=np.int32)
                batch_gt_boxes = np.zeros(
                    (batch_size, config.MAX_GT_INSTANCES, 4), dtype=np.int32)
                batch_gt_masks = np.zeros(
                    (batch_size, gt_masks.shape[0], gt_masks.shape[1],
                     config.MAX_GT_INSTANCES), dtype=gt_masks.dtype)
                if random_rois:
                    batch_rpn_rois = np.zeros(
                        (batch_size, rpn_rois.shape[0], 4), dtype=rpn_rois.dtype)
                    if detection_targets:
                        batch_rois = np.zeros(
                            (batch_size,) + rois.shape, dtype=rois.dtype)
                        batch_mrcnn_class_ids = np.zeros(
                            (batch_size,) + mrcnn_class_ids.shape, dtype=mrcnn_class_ids.dtype)
                        batch_mrcnn_bbox = np.zeros(
                            (batch_size,) + mrcnn_bbox.shape, dtype=mrcnn_bbox.dtype)
                        batch_mrcnn_mask = np.zeros(
                            (batch_size,) + mrcnn_mask.shape, dtype=mrcnn_mask.dtype)

            # If more instances than fits in the array, sub-sample from them.
            if gt_boxes.shape[0] > config.MAX_GT_INSTANCES:
                ids = np.random.choice(
                    np.arange(gt_boxes.shape[0]), config.MAX_GT_INSTANCES, replace=False)
                gt_class_ids = gt_class_ids[ids]
                gt_boxes = gt_boxes[ids]
                gt_masks = gt_masks[:, :, ids]

            # Add to batch
            batch_image_meta[b] = image_meta
            batch_rpn_match[b] = rpn_match[:, np.newaxis]
            batch_rpn_bbox[b] = rpn_bbox
            batch_images[b] = mold_image(image.astype(np.float32), config)
            batch_gt_class_ids[b, :gt_class_ids.shape[0]] = gt_class_ids
            batch_gt_boxes[b, :gt_boxes.shape[0]] = gt_boxes
            batch_gt_masks[b, :, :, :gt_masks.shape[-1]] = gt_masks
            if random_rois:
                batch_rpn_rois[b] = rpn_rois
                if detection_targets:
                    batch_rois[b] = rois
                    batch_mrcnn_class_ids[b] = mrcnn_class_ids
                    batch_mrcnn_bbox[b] = mrcnn_bbox
                    batch_mrcnn_mask[b] = mrcnn_mask
            b += 1

            # Batch full?
            if b >= batch_size:
                inputs = [batch_images, batch_image_meta, batch_rpn_match, batch_rpn_bbox,
                          batch_gt_class_ids, batch_gt_boxes, batch_gt_masks]
                outputs = []

                if random_rois:
                    inputs.extend([batch_rpn_rois])
                    if detection_targets:
                        inputs.extend([batch_rois])
                        # Keras requires that output and targets have the same number of dimensions
                        batch_mrcnn_class_ids = np.expand_dims(
                            batch_mrcnn_class_ids, -1)
                        outputs.extend(
                            [batch_mrcnn_class_ids, batch_mrcnn_bbox, batch_mrcnn_mask])

                yield inputs, outputs

                # start a new batch
                b = 0
        except (GeneratorExit, KeyboardInterrupt):
            raise
        except:
            # Log it and skip the image
            logging.exception("Error processing image {}".format(
                dataset.image_info[image_id]))
            error_count += 1
            if error_count > 5:
                raise


############################################################
#  MaskRCNN Class
############################################################
class MaskRCNN(object):
    def __init__(self, mode, config, model_dir):
        assert mode in ['training', 'inference']
        self.mode = mode
        self.config = config
        self.model_dir = model_dir
        self.model_input, self.model_output = self.build(mode=mode, config=config)
        self.summaries = []
        sessconfig = tf.ConfigProto(allow_soft_placement=True)
        sessconfig.gpu_options.allow_growth = True
        self.sess = tf.Session(config=sessconfig)
        self.saver = tf.train.Saver(max_to_keep=1)
        # self.summary_writer = tf.summary.FileWriter('./logs/', self.sess.graph)

    def build(self, mode, config):
        """Build Mask R-CNN architecture.
                    input_shape: The shape of the input image.
                    mode: Either "training" or "inference". The inputs and
                        outputs of the model differ accordingly.
        """
        assert mode in ['training', 'inference']
        regularizer = tf.contrib.layers.l2_regularizer(self.config.WEIGHT_DECAY)
        # regularizer = None
        # Image size must be dividable by 2 multiple times
        h, w = config.IMAGE_SHAPE[:2]
        if h / 2 ** 6 != int(h / 2 ** 6) or w / 2 ** 6 != int(w / 2 ** 6):
            raise Exception("Image size must be dividable by 2 at least 6 times "
                            "to avoid fractions when downscaling and upscaling."
                            "For example, use 256, 320, 384, 448, 512, ... etc. ")
        # Inputs
        input_image = tf.placeholder(shape=[None, None, None, config.IMAGE_SHAPE[2]], name="input_image", dtype=tf.float32)
        input_image_meta = tf.placeholder(shape=[None, config.IMAGE_META_SIZE], name="input_image_meta", dtype=tf.float32)

        if mode == "training":
            # RPN GT
            input_rpn_match = tf.placeholder(shape=[None, None, 1], name="input_rpn_match", dtype=tf.int32)
            # ground truth的一幅图像有几个框，每个框的分类
            input_rpn_bbox = tf.placeholder(shape=[None, None, 4], name="input_rpn_bbox", dtype=tf.float32)
            # ground truth的一幅图像有几个框，每个框的坐标

            # Detection GT (class IDs, bounding boxes, and masks)
            # 1. GT Class IDs (zero padded)
            # ground truth的一幅图像有几个框，每个框的分类
            input_gt_class_ids = tf.placeholder(shape=[None, None], name="input_gt_class_ids", dtype=tf.int32)
            # 2. GT Boxes in pixels (zero padded)
            # [batch, MAX_GT_INSTANCES, (y1, x1, y2, x2)] in image coordinates
            # ground truth的一幅图像有几个框，每个框的坐标，
            # 从上面的注释看，这个框的坐标是两个对角坐标，而不是coco原本的左上角坐标加长宽，是不是他自己变过了
            input_gt_boxes = tf.placeholder(shape=[None, None, 4], name="input_gt_boxes", dtype=tf.float32)
            # Normalize coordinates
            # 这里算shape的时候把batch也算进去了，所以是1:3，因为0是batch。
            # 把输入的坐标归一化，从源码上看，并不是把图像中心作为0，而是把图像左上角，即原来的(0,0)位置依旧是(0,0)
            gt_boxes = norm_boxes_graph(input_gt_boxes, tf.shape(input_image)[1:3])
            # 3. GT Masks (zero padded)
            # [batch, height, width, MAX_GT_INSTANCES]
            # 输入的mask得是zero padding的
            if config.USE_MINI_MASK:
                input_gt_masks = tf.placeholder(
                    shape=[None, config.MINI_MASK_SHAPE[0], config.MINI_MASK_SHAPE[1], None],
                    name="input_gt_masks", dtype=bool)
            else:
                input_gt_masks = tf.placeholder(
                    shape=[None, config.IMAGE_SHAPE[0], config.IMAGE_SHAPE[1], None],
                    name="input_gt_masks", dtype=bool)
        elif mode == "inference":
            # Anchors in normalized coordinates
            input_anchors = tf.placeholder(shape=[None, None, 4], name="input_anchors", dtype=tf.float32)

        # Build the shared convolutional layers.
        # Bottom-up Layers
        # Returns a list of the last layers of each stage, 5 in total.
        # Don't create the thead (stage 5), so we pick the 4th item in the list.
        if callable(config.BACKBONE):
            _, C2, C3, C4, C5 = config.BACKBONE(input_image, stage5=True,
                                                train_bn=config.TRAIN_BN)
        else:
            _, C2, C3, C4, C5 = resnet_graph(input_image, config.BACKBONE,
                                             stage5=True, train_bn=config.TRAIN_BN, use_bias=config.USE_BIAS, regularizer=regularizer)
        P5 = tf.layers.conv2d(C5, filters=config.TOP_DOWN_PYRAMID_SIZE, kernel_size=(1, 1),
                              name='fpn_c5p5', kernel_regularizer=regularizer)
        P4 = tf.add(
            tf.keras.layers.UpSampling2D(size=(2, 2), name="fpn_p5upsampled")(P5),
            tf.layers.conv2d(C4, config.TOP_DOWN_PYRAMID_SIZE, (1, 1), name='fpn_c4p4', kernel_regularizer=regularizer),
            name="fpn_p4add")
        P3 = tf.add(
            tf.keras.layers.UpSampling2D(size=(2, 2), name="fpn_p4upsampled")(P4),
            tf.layers.conv2d(C3, config.TOP_DOWN_PYRAMID_SIZE, (1, 1), name='fpn_c3p3', kernel_regularizer=regularizer),
            name="fpn_p3add")
        P2 = tf.add(
            tf.keras.layers.UpSampling2D(size=(2, 2), name="fpn_p3upsampled")(P3),
            tf.layers.conv2d(C2, config.TOP_DOWN_PYRAMID_SIZE, (1, 1), name='fpn_c2p2', kernel_regularizer=regularizer),
            name="fpn_p2add")
        # Attach 3x3 conv to all P layers to get the final feature maps.
        P2 = tf.layers.conv2d(P2, config.TOP_DOWN_PYRAMID_SIZE, (3, 3), padding="SAME", name="fpn_p2", kernel_regularizer=regularizer)
        P3 = tf.layers.conv2d(P3, config.TOP_DOWN_PYRAMID_SIZE, (3, 3), padding="SAME", name="fpn_p3", kernel_regularizer=regularizer)
        P4 = tf.layers.conv2d(P4, config.TOP_DOWN_PYRAMID_SIZE, (3, 3), padding="SAME", name="fpn_p4", kernel_regularizer=regularizer)
        P5 = tf.layers.conv2d(P5, config.TOP_DOWN_PYRAMID_SIZE, (3, 3), padding="SAME", name="fpn_p5", kernel_regularizer=regularizer)
        # P6 is used for the 5th anchor scale in RPN. Generated by
        # subsampling from P5 with stride of 2.
        P6 = tf.layers.max_pooling2d(P5, pool_size=(1, 1), strides=2, name="fpn_p6")

        # Note that P6 is used in RPN, but not in the classifier heads.
        rpn_feature_maps = [P2, P3, P4, P5, P6]  # 对应的anchor scale分别是[32, 64, 128, 256, 512]
        mrcnn_feature_maps = [P2, P3, P4, P5]

        if mode == "training":
            # 返回归一化之后的anchor box坐标，这个anchor box的大小是和feature map的level有关的。
            # size=[anchor_count, 4], 当然这里面的anchor box大小不一，是从小到大的。
            # 还有这个anchor box的坐标是在输入图像上的坐标
            anchors = self.get_anchors(config.IMAGE_SHAPE)  # 这个IMAGE_SHAPE应该是1024*1024
            # Duplicate across the batch dimension because Keras requires it
            # TODO: can this be optimized to avoid duplicating the anchors?
            # 实际上，因为输入图像的大小固定为[1024, 1024], 所以batch中每一个图片的anchor都是一样的，这里还要把它broadcast
            anchors = np.broadcast_to(anchors, (config.BATCH_SIZE,) + anchors.shape)
            anchors = tf.constant(anchors, name="anchors")
        else:
            anchors = input_anchors
            # 当inference的时候，输入input_anchors
            # 这其实挺奇怪的，难道inference的时候anchor不是用feature map产生的吗？
            # 难道是提前计算好，反正anchor的位置只与feature map的shape、输入图像的shape有关，和起具体参数是多少无关，ratio也是知道的
            # 完全可以提前算好，节省一点时间？

        # RPN Model
        # Loop through pyramid layers
        layer_outputs = []  # list of lists
        reuse = False
        for p in rpn_feature_maps:
            layer_outputs.append(rpn_model(p, config.RPN_ANCHOR_STRIDE,
                                           len(config.RPN_ANCHOR_RATIOS), reuse, regularizer=regularizer))
            reuse = True
        # Concatenate layer outputs
        # Convert from list of lists of level outputs to list of lists
        # of outputs across levels.
        # e.g. [[a1, b1, c1], [a2, b2, c2]] => [[a1, a2], [b1, b2], [c1, c2]]
        output_names = ["rpn_class_logits", "rpn_class", "rpn_bbox"]
        outputs = list(zip(*layer_outputs))
        outputs = [tf.concat(list(o), axis=1, name=n) for o, n in zip(outputs, output_names)]

        rpn_class_logits, rpn_class, rpn_bbox = outputs
        # rpn_class_logits: [batch, anchor_N, 2]
        # rpn_class: [batch, anchor_N, 2]
        # rpn_bbox: [batch, anchor_N, 4]

        # Generate proposals
        # Proposals are [batch, proposal_count, (y1, x1, y2, x2)] in normalized coordinates and zero padded.
        proposal_count = config.POST_NMS_ROIS_TRAINING if mode == "training" else config.POST_NMS_ROIS_INFERENCE
        # proposal_count=2000 if training else 1000
        rpn_rois = ProposalLayer(proposal_count=proposal_count,
                                 nms_threshold=config.RPN_NMS_THRESHOLD,
                                 name="ROI",
                                 config=config)([rpn_class, rpn_bbox, anchors])
        # 接受rpn产生的scores、deltas和anchor，用scores选取一定数量的anchor，再用deltas对anchor进行调整
        # 再用非极大抑制对anchor进行筛选,输出感兴区域

        if mode == "training":
            # Class ID mask to mark class IDs supported by the dataset the image
            # came from.
            active_class_ids = parse_image_meta_graph(input_image_meta)["active_class_ids"]
            # 输入input_image_meta，输出input_image_meta中包含的属性，是一个字典。

            if not config.USE_RPN_ROIS:  # USE_RPN_ROIS=True
                # Ignore predicted ROIs and use ROIs provided as an input.
                input_rois = tf.placeholder(shape=[None, config.POST_NMS_ROIS_TRAINING, 4],
                                            name="input_roi", dtype=np.int32)
                # 如果不使用rpn产生的ROI，那么就要从外界输入ROI。
                # Normalize coordinates
                target_rois = norm_boxes_graph(input_rois, tf.shape(input_image)[1:3])
            else:
                target_rois = rpn_rois

            # Generate detection targets
            # Subsamples proposals and generates target outputs for training
            # Note that proposal class IDs, gt_boxes, and gt_masks are zero
            # padded. Equally, returned rois and targets are zero padded.

            rois, target_class_ids, target_bbox, target_mask = \
                DetectionTargetLayer(config, name="proposal_targets")([target_rois, input_gt_class_ids,
                                                                       gt_boxes, input_gt_masks])
            self.Detectiontarget_output = [rois, target_class_ids, target_bbox, target_mask]

            # ind = np.load('ind.npy').astype(np.int32)
            # ind = np.reshape(ind, (ind.shape[0], 1))
            # gather_ind = np.concatenate((np.zeros((ind.shape[0], 1)).astype(np.int32), ind), -1)
            # rois = tf.gather_nd(rpn_rois, gather_ind)
            # rois = tf.expand_dims(rois, 0)

            # 输入RPN的proposal，gt_class_ids, gt_boxes, gt_masks。
            # 从proposal中提取出200个ROI，其中包含positive ROI和negative ROI，如果两者加起来还不足200，就用0padding，再输出
            # 同时也要输出positive ROI对应的gt_class_ids, bbox_refine_delta(上面的就是gt_boxes), gt_masks
            # 再输入到network head中

            # Network Heads
            # TODO: verify that this handles zero padded ROIs
            mrcnn_class_logits, mrcnn_class, mrcnn_bbox = \
                fpn_classifier_graph(rois, mrcnn_feature_maps, input_image_meta,
                                     config.POOL_SIZE, config.NUM_CLASSES,
                                     train_bn=config.TRAIN_BN,
                                     fc_layers_size=config.FPN_CLASSIF_FC_LAYERS_SIZE,
                                     regularizer=regularizer)

            mrcnn_mask = build_fpn_mask_graph(rois, mrcnn_feature_maps,
                                              input_image_meta,
                                              config.MASK_POOL_SIZE,
                                              config.NUM_CLASSES,
                                              train_bn=config.TRAIN_BN,
                                              regularizer=regularizer)

            # TODO: clean up (use tf.identify if necessary)
            output_rois = tf.identity(rois, name="output_rois")

            # Losses
            # rpn_class_loss = 0.0
            # rpn_bbox_loss = 0.0
            # class_loss = 0.0
            # bbox_loss = 0.0
            # mask_loss = 0.0

            rpn_class_loss = rpn_class_loss_graph(input_rpn_match, rpn_class_logits)
            rpn_bbox_loss = rpn_bbox_loss_graph(config, input_rpn_bbox, input_rpn_match, rpn_bbox)
            class_loss = mrcnn_class_loss_graph(target_class_ids, mrcnn_class_logits, active_class_ids)
            bbox_loss = mrcnn_bbox_loss_graph(target_bbox, target_class_ids, mrcnn_bbox)
            mask_loss = mrcnn_mask_loss_graph(target_mask, target_class_ids, mrcnn_mask)

            rpn_class_loss = tf.identity(rpn_class_loss, name="rpn_class_loss")
            rpn_bbox_loss = tf.identity(rpn_bbox_loss, name="rpn_bbox_loss")
            class_loss = tf.identity(class_loss, name="mrcnn_class_loss")
            bbox_loss = tf.identity(bbox_loss, name="mrcnn_bbox_loss")
            mask_loss = tf.identity(mask_loss, name="mrcnn_mask_loss")

            inputs = [input_image, input_image_meta,
                      input_rpn_match, input_rpn_bbox,
                      input_gt_boxes, input_gt_class_ids, input_gt_masks]
            outputs = [rpn_class_logits, rpn_class, rpn_bbox,
                       mrcnn_class_logits, mrcnn_class, mrcnn_bbox, mrcnn_mask,
                       rpn_rois, output_rois,
                       rpn_class_loss, rpn_bbox_loss, class_loss, bbox_loss, mask_loss]
        else:
            # ind = np.load('ind.npy').astype(np.int32)
            # ind = np.reshape(ind, (ind.shape[0], 1))
            # gather_ind = np.concatenate((np.zeros((ind.shape[0], 1)).astype(np.int32), ind), -1)
            # rpn_rois = tf.gather_nd(rpn_rois, gather_ind)
            # rpn_rois = tf.expand_dims(rpn_rois, 0)
            # Network Heads
            # Proposal classifier and BBox regressor heads
            mrcnn_class_logits, mrcnn_class, mrcnn_bbox, =\
                fpn_classifier_graph(rpn_rois, mrcnn_feature_maps, input_image_meta,
                                     config.POOL_SIZE, config.NUM_CLASSES,
                                     train_bn=config.TRAIN_BN,
                                     fc_layers_size=config.FPN_CLASSIF_FC_LAYERS_SIZE,
                                     regularizer=regularizer)
            # Detections
            # output is [batch, num_detections, (y1, x1, y2, x2, class_id, score)] in
            # normalized coordinates
            detections = DetectionLayer(config, name="mrcnn_detection")([rpn_rois, mrcnn_class,
                                                                         mrcnn_bbox, input_image_meta])
            # 利用classifier生成的box和class确定要生成哪个类的mask，每一个box都对应一个类，即上面的class_id。

            # Create masks for detections
            detection_boxes = detections[..., :4]  # [batch, num_detections, (y1, x1, y2, x2)]
            mrcnn_mask = build_fpn_mask_graph(detection_boxes, mrcnn_feature_maps,
                                              input_image_meta,
                                              config.MASK_POOL_SIZE,
                                              config.NUM_CLASSES,
                                              train_bn=config.TRAIN_BN,
                                              regularizer=regularizer)

            inputs = [input_image, input_image_meta, input_anchors]
            outputs = [detections, mrcnn_class_logits, mrcnn_class, mrcnn_bbox,
                       mrcnn_mask, rpn_rois, rpn_class, rpn_bbox]

        return inputs, outputs

    def get_anchors(self, image_shape):
        """Returns anchor pyramid for the given image size."""
        backbone_shapes = compute_backbone_shapes(self.config, image_shape)
        # 返回应该是[[256,256],[128,128],[64, 64],[32, 32],[16, 16]]
        # Cache anchors and reuse if image shape is the same
        if not hasattr(self, "_anchor_cache"):
            self._anchor_cache = {}  # 这是一个字典，关键字是tuple(image_shape)
        if not tuple(image_shape) in self._anchor_cache:
            # Generate Anchors
            # 返回每一个feature map(从大到小)的anchor box，从下面的代码上看顺序就是这样的，# size=[anchor_count, 4]
            # 这里还没有归一化
            a = utils.generate_pyramid_anchors(
                self.config.RPN_ANCHOR_SCALES,  # (32, 64, 128, 256, 512)  不知道为什么这个要用元组
                self.config.RPN_ANCHOR_RATIOS,  # [0.5, 1, 2]
                backbone_shapes,  # [[256,256],[128,128],[64, 64],[32, 32],[16, 16]]
                self.config.BACKBONE_STRIDES,  # [4, 8, 16, 32, 64]
                self.config.RPN_ANCHOR_STRIDE)  # 1
            # Keep a copy of the latest anchors in pixel coordinates because
            # it's used in inspect_model notebooks.
            # TODO: Remove this after the notebook are refactored to not use it
            self.anchors = a
            # Normalize coordinates
            self._anchor_cache[tuple(image_shape)] = utils.norm_boxes(a, image_shape[:2])
        return self._anchor_cache[tuple(image_shape)]  # 返回归一化之后的anchor box坐标

    def train(self, train_dataset, val_dataset, learning_rate, epochs, layers,
              augmentation=None, custom_callbacks=None, no_augmentation_sources=None):

        assert self.mode == "training", "Create model in training mode."

        # Data generators
        train_generator = data_generator(train_dataset, self.config, shuffle=True,
                                         augmentation=augmentation,
                                         batch_size=self.config.BATCH_SIZE,
                                         no_augmentation_sources=no_augmentation_sources)
        # train_generator = data_generator(train_dataset, self.config, shuffle=False,
        #                                  batch_size=self.config.BATCH_SIZE)
        val_generator = data_generator(val_dataset, self.config, shuffle=True,
                                       batch_size=self.config.BATCH_SIZE)

        self.optimize(layers, learning_rate)
        self.sess.run(tf.global_variables_initializer())

        self.summary_op = tf.summary.merge(self.summaries)
        self.summary_writer = tf.summary.FileWriter('./logs/', self.sess.graph)

        # assign_ops = self.checkpoint_assign(path='./rpn_dic1.npy')
        # self.sess.run(assign_ops)

        # assign_ops = self.imagenet_assign()
        # self.sess.run(assign_ops)

        # 抽出几个variable看看assign是否正常
        # var = []
        # for v in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES):
        #     if v.name.startswith('res2a_branch1'):
        #         var.append(v)
        # var_np = self.sess.run(var)

        # load pretrained checkpoint model
        self.saver.restore(self.sess, './checkpoint/model-319')

        single_train_inputs, _ = train_generator.__next__()
        single_val_inputs, _ = val_generator.__next__()

        start = 320
        for epoch in range(start, start+epochs):
            print("Epoch {}/{}.".format(epoch, epochs))
            self.train_epoch(train_generator, epoch, single_inputs=None)
            self.val_epoch(val_generator, epoch, single_inputs=None)
            checkpoint_path = './checkpoint/model'
            self.saver.save(self.sess, checkpoint_path, global_step=epoch)

    def train_epoch(self, DataGenerator, epoch, single_inputs=None):
        steps = self.config.STEPS_PER_EPOCH
        step = 0
        rpn_class_loss = 0
        rpn_bbox_loss = 0
        class_loss = 0
        bbox_loss = 0
        mask_loss = 0

        for inputs, outputs in DataGenerator:
            images = inputs[0]
            image_metas = inputs[1]
            rpn_match = inputs[2]
            rpn_bbox = inputs[3]
            gt_class_ids = inputs[4]
            gt_boxes = inputs[5]
            gt_masks = inputs[6]

            if single_inputs:
                images = single_inputs[0]
                image_metas = single_inputs[1]
                rpn_match = single_inputs[2]
                rpn_bbox = single_inputs[3]
                gt_class_ids = single_inputs[4]
                gt_boxes = single_inputs[5]
                gt_masks = single_inputs[6]

            input_image = self.model_input[0]
            input_image_meta = self.model_input[1]
            input_rpn_match = self.model_input[2]
            input_rpn_bbox = self.model_input[3]
            input_gt_boxes = self.model_input[4]
            input_gt_class_ids = self.model_input[5]
            input_gt_masks = self.model_input[6]

            t1 = time.time()
            rpn_output, mrcnn_output, proposal_detectiontarget_roi, loss_list, DetectTarget_output, _ = \
                self.sess.run([self.rpn_output, self.mrcnn_output,
                               self.Proposal_and_DetectionTarget_ROI, self.loss_list,
                               self.Detectiontarget_output,
                               self.train_op],
                              feed_dict={input_image: images,
                                         input_image_meta: image_metas,
                                         input_rpn_match: rpn_match,
                                         input_rpn_bbox: rpn_bbox,
                                         input_gt_boxes: gt_boxes,
                                         input_gt_class_ids: gt_class_ids,
                                         input_gt_masks: gt_masks})

            assert not math.isnan(sum(loss_list)), "nan, stop"

            t2 = time.time()
            remain_time = (t2-t1) * ((320 - epoch) * steps + (steps - step)) / (60 * 60 * 24)

            sys.stdout.write('\r{}/{}; '
                             '[rpn_class_loss: {:.4}, rpn_bbox_loss: {:.4}, class_loss: {:.4}, bbox_loss: {:.4}, mask_loss: {:.4}, '
                             'loss_sum: {:.4}, l2_loss: {:.4}, loss_total: {:.4}]'
                             '[c_time/re_time: {:.2}/{:.2} days]'.
                             format(step, steps,
                                    loss_list[0], loss_list[1], loss_list[2], loss_list[3], loss_list[4],
                                    loss_list[5], loss_list[6], loss_list[7],
                                    (t2 - t1), remain_time))

            # statistic
            rpn_class_loss += loss_list[0] / steps
            rpn_bbox_loss += loss_list[1] / steps
            class_loss += loss_list[2] / steps
            bbox_loss += loss_list[3] / steps
            mask_loss += loss_list[4] / steps

            if step % 100 == 0:
                summary_str = self.sess.run(self.summary_op,
                                            feed_dict={input_image: images,
                                                       input_image_meta: image_metas,
                                                       input_rpn_match: rpn_match,
                                                       input_rpn_bbox: rpn_bbox,
                                                       input_gt_boxes: gt_boxes,
                                                       input_gt_class_ids: gt_class_ids,
                                                       input_gt_masks: gt_masks})
                self.summary_writer.add_summary(summary_str, step + epoch*steps)
                # checkpoint_path = './checkpoint/model'
                # self.saver.save(self.sess, checkpoint_path, global_step=epoch)
            if step == (steps - 1):
                break

            step += 1

        print('\naverage_loss: '
              '[rpn_class_loss: {:.4}, rpn_bbox_loss: {:.4}, class_loss: {:.4}, bbox_loss: {:.4}, mask_loss: {:.4}]'.
              format(rpn_class_loss, rpn_bbox_loss, class_loss, bbox_loss, mask_loss))

    def val_epoch(self, DataGenerator, epoch, single_inputs=None):
        steps = self.config.VALIDATION_STEPS
        step = 0
        rpn_class_loss = 0
        rpn_bbox_loss = 0
        class_loss = 0
        bbox_loss = 0
        mask_loss = 0

        for inputs, outputs in DataGenerator:
            images = inputs[0]
            image_metas = inputs[1]
            rpn_match = inputs[2]
            rpn_bbox = inputs[3]
            gt_class_ids = inputs[4]
            gt_boxes = inputs[5]
            gt_masks = inputs[6]

            if single_inputs:
                images = single_inputs[0]
                image_metas = single_inputs[1]
                rpn_match = single_inputs[2]
                rpn_bbox = single_inputs[3]
                gt_class_ids = single_inputs[4]
                gt_boxes = single_inputs[5]
                gt_masks = single_inputs[6]

            input_image = self.model_input[0]
            input_image_meta = self.model_input[1]
            input_rpn_match = self.model_input[2]
            input_rpn_bbox = self.model_input[3]
            input_gt_boxes = self.model_input[4]
            input_gt_class_ids = self.model_input[5]
            input_gt_masks = self.model_input[6]

            loss_list = \
                self.sess.run(
                    self.loss_list,
                    feed_dict={input_image: images,
                               input_image_meta: image_metas,
                               input_rpn_match: rpn_match,
                               input_rpn_bbox: rpn_bbox,
                               input_gt_boxes: gt_boxes,
                               input_gt_class_ids: gt_class_ids,
                               input_gt_masks: gt_masks})

            sys.stdout.write('\rstep/steps:{}/{}; '
                             '[rpn_class_loss: {:.4}, rpn_bbox_loss: {:.4}, class_loss: {:.4}, bbox_loss: {:.4}, mask_loss: {:.4}]'.
                             format(step, steps, loss_list[0], loss_list[1], loss_list[2], loss_list[3], loss_list[4]))

            # statistic
            rpn_class_loss += loss_list[0] / steps
            rpn_bbox_loss += loss_list[1] / steps
            class_loss += loss_list[2] / steps
            bbox_loss += loss_list[3] / steps
            mask_loss += loss_list[4] / steps

            if step == (steps - 1):
                break

            step += 1

        print('\naverage_loss: '
              '[rpn_class_loss: {:.4}, rpn_bbox_loss: {:.4}, class_loss: {:.4}, bbox_loss: {:.4}, mask_loss: {:.4}]'.
              format(rpn_class_loss, rpn_bbox_loss, class_loss, bbox_loss, mask_loss))

    def optimize(self, layers, learning_rate):
        """
        layers: Allows selecting wich layers to train. It can be:
            - A regular expression to match layer names to train
            - One of these predefined values:
              heads: The RPN, classifier and mask heads of the network
              all: All the layers
              3+: Train Resnet stage 3 and up
              4+: Train Resnet stage 4 and up
              5+: Train Resnet stage 5 and up
        """
        layer_regex = {
            # all layers but the backbone
            "heads": r"(mrcnn\_.*)|(rpn\_.*)|(fpn\_.*)",
            # From a specific Resnet stage and up
            "3+": r"(scale3.*)|(scale4.*)|(scale5.*)|(mrcnn\_.*)|(rpn\_.*)|(fpn\_.*)",
            "4+": r"(scale4.*)|(scale5.*)|(mrcnn\_.*)|(rpn\_.*)|(fpn\_.*)",
            "5+": r"(scale5.*)|(mrcnn\_.*)|(rpn\_.*)|(fpn\_.*)",
            # All layers
            "all": ".*",
            # only rpn
            "rpn": r"(rpn\_.*)|(fpn\_.*)",
            "mrcnn": r"(mrcnn\_.*)"
        }
        layer_flag = layers
        if layers in layer_regex.keys():
            layers = layer_regex[layers]
        var_need_train = self.set_trainbale(layers_regex=layers)

        rpn_class_logits, rpn_class, rpn_bbox = self.model_output[0:3]
        self.rpn_output = [rpn_class_logits, rpn_class, rpn_bbox]
        mrcnn_class_logits, mrcnn_class, mrcnn_bbox, mrcnn_mask = self.model_output[3:7]
        self.mrcnn_output = [mrcnn_class_logits, mrcnn_class, mrcnn_bbox, mrcnn_mask]
        rpn_rois, output_rois = self.model_output[7:9]
        self.Proposal_and_DetectionTarget_ROI = [rpn_rois, output_rois]
        rpn_class_loss, rpn_bbox_loss, class_loss, bbox_loss, mask_loss = self.model_output[9:]
        self.loss_list = [rpn_class_loss, rpn_bbox_loss, class_loss, bbox_loss, mask_loss]

        if layer_flag == "rpn":
            loss_sum = tf.reduce_sum(self.loss_list[:2])  # 只训练rpn网络
        elif layer_flag == "mrcnn":
            loss_sum = tf.reduce_sum(self.loss_list[3:])  # 只训练mask head
        else:
            loss_weight = tf.constant([0.1, 0.1, 0.6, 0.1, 0.1], dtype=tf.float32)
            loss_sum = tf.reduce_sum(tf.stack(self.loss_list) * loss_weight)
            # loss_sum = tf.reduce_sum(self.loss_list)
        self.loss_list.append(loss_sum)
        self.summaries.append(tf.summary.scalar('loss_sum', loss_sum))
        self.summaries.append(tf.summary.scalar('rpn/class_loss', rpn_class_loss))
        self.summaries.append(tf.summary.scalar('rpn/bbox_loss', rpn_bbox_loss))
        self.summaries.append(tf.summary.scalar('class_loss', class_loss))
        self.summaries.append(tf.summary.scalar('bbox_loss', bbox_loss))
        self.summaries.append(tf.summary.scalar('mask_loss', mask_loss))
        l2_loss = tf.losses.get_regularization_loss()
        self.summaries.append(tf.summary.scalar('l2_loss', l2_loss))
        self.loss_list.append(l2_loss)
        loss_total = loss_sum + l2_loss
        self.loss_list.append(loss_total)
        global_step = tf.Variable(0, trainable=False)
        start_learning_rate = learning_rate
        # learning_rate_decay = tf.train.exponential_decay(learning_rate=start_learning_rate, global_step=global_step,
        #                                            decay_steps=160000, decay_rate=0.5,
        #                                            staircase=True)
        learning_rate_decay = 0.000025
        self.summaries.append(tf.summary.scalar('learning_rate', learning_rate_decay))
        optim = tf.train.AdamOptimizer(learning_rate=learning_rate_decay, beta2=self.config.LEARNING_MOMENTUM)
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        with tf.control_dependencies(update_ops):
            self.apply_grad_op = optim.minimize(loss_total, global_step=global_step, var_list=var_need_train)
            # if consider weight decay: use loss_total; if not, use loss_sum
        ema = tf.train.ExponentialMovingAverage(decay=0.99, num_updates=global_step)
        with tf.control_dependencies([self.apply_grad_op]):
            self.variable_average_op = ema.apply(tf.trainable_variables())

        self.train_op = tf.group(self.apply_grad_op, self.variable_average_op)

    def imagenet_assign(self):
        var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
        weight = np.load('resnet_weight_new_name.npy')[0]
        assign_op = []
        num = 0
        for var in var_list:
            var_name = var.name
            if var_name in weight:
                assign_op.append(tf.assign(var, weight[var_name]))
                # print('assign variable {}'.format(var_name))
                num += 1
            # else:
            #     print('error: there is no variable in weights named {}'.format(var_name))

        return assign_op

    def checkpoint_assign(self, path):
        var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
        weight = np.load(path)[0]
        assign_op = []
        num = 0
        for var in var_list:
            var_name = var.name
            if var_name in weight:
                assign_op.append(tf.assign(var, weight[var_name]))
                print('assign variable {}'.format(var_name))
                num += 1
            else:
                print('error: there is no variable in weights named {}'.format(var_name))

        return assign_op

    def set_trainbale(self, layers_regex):
        var_need_train = []
        for var in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES):
            var_name = var.name
            need_train = bool(re.fullmatch(layers_regex, var_name))
            if need_train:
                var_need_train.append(var)

        return var_need_train

    def mold_inputs(self, images):
        """Takes a list of images and modifies them to the format expected
        as an input to the neural network.
        images: List of image matrices [height,width,depth]. Images can have
            different sizes.

        Returns 3 Numpy matrices:
        molded_images: [N, h, w, 3]. Images resized and normalized.
        image_metas: [N, length of meta data]. Details about each image.
        windows: [N, (y1, x1, y2, x2)]. The portion of the image that has the
            original image (padding excluded).
        """
        molded_images = []
        image_metas = []
        windows = []
        for image in images:
            # Resize image
            # TODO: move resizing to mold_image()
            molded_image, window, scale, padding, crop = utils.resize_image(
                image,
                min_dim=self.config.IMAGE_MIN_DIM,
                min_scale=self.config.IMAGE_MIN_SCALE,
                max_dim=self.config.IMAGE_MAX_DIM,
                mode=self.config.IMAGE_RESIZE_MODE)
            molded_image = mold_image(molded_image, self.config)
            # Build image_meta
            image_meta = compose_image_meta(
                0, image.shape, molded_image.shape, window, scale,
                np.zeros([self.config.NUM_CLASSES], dtype=np.int32))
            # Append
            molded_images.append(molded_image)
            windows.append(window)
            image_metas.append(image_meta)
        # Pack into arrays
        molded_images = np.stack(molded_images)
        image_metas = np.stack(image_metas)
        windows = np.stack(windows)
        return molded_images, image_metas, windows

    def unmold_detections(self, detections, mrcnn_mask, original_image_shape,
                          image_shape, window):
        """Reformats the detections of one image from the format of the neural
        network output to a format suitable for use in the rest of the
        application.

        detections: [N, (y1, x1, y2, x2, class_id, score)] in normalized coordinates
        mrcnn_mask: [N, height, width, num_classes]
        original_image_shape: [H, W, C] Original image shape before resizing
        image_shape: [H, W, C] Shape of the image after resizing and padding
        window: [y1, x1, y2, x2] Pixel coordinates of box in the image where the real
                image is excluding the padding.

        Returns:
        boxes: [N, (y1, x1, y2, x2)] Bounding boxes in pixels
        class_ids: [N] Integer class IDs for each bounding box
        scores: [N] Float probability scores of the class_id
        masks: [height, width, num_instances] Instance masks
        """
        # How many detections do we have?
        # Detections array is padded with zeros. Find the first class_id == 0.
        zero_ix = np.where(detections[:, 4] == 0)[0]
        N = zero_ix[0] if zero_ix.shape[0] > 0 else detections.shape[0]

        # Extract boxes, class_ids, scores, and class-specific masks
        boxes = detections[:N, :4]
        class_ids = detections[:N, 4].astype(np.int32)
        scores = detections[:N, 5]
        masks = mrcnn_mask[np.arange(N), :, :, class_ids]

        # Translate normalized coordinates in the resized image to pixel
        # coordinates in the original image before resizing
        window = utils.norm_boxes(window, image_shape[:2])
        wy1, wx1, wy2, wx2 = window
        shift = np.array([wy1, wx1, wy1, wx1])
        wh = wy2 - wy1  # window height
        ww = wx2 - wx1  # window width
        scale = np.array([wh, ww, wh, ww])
        # Convert boxes to normalized coordinates on the window
        boxes = np.divide(boxes - shift, scale)
        # Convert boxes to pixel coordinates on the original image
        boxes = utils.denorm_boxes(boxes, original_image_shape[:2])

        # Filter out detections with zero area. Happens in early training when
        # network weights are still random
        exclude_ix = np.where(
            (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) <= 0)[0]
        if exclude_ix.shape[0] > 0:
            boxes = np.delete(boxes, exclude_ix, axis=0)
            class_ids = np.delete(class_ids, exclude_ix, axis=0)
            scores = np.delete(scores, exclude_ix, axis=0)
            masks = np.delete(masks, exclude_ix, axis=0)
            N = class_ids.shape[0]

        # Resize masks to original image size and set boundary threshold.
        full_masks = []
        for i in range(N):
            # Convert neural network mask to full size mask
            full_mask = utils.unmold_mask(masks[i], boxes[i], original_image_shape)
            full_masks.append(full_mask)
        full_masks = np.stack(full_masks, axis=-1)\
            if full_masks else np.empty(original_image_shape[:2] + (0,))

        return boxes, class_ids, scores, full_masks

    def detect(self, images):
        """Runs the detection pipeline.

        images: List of images, potentially of different sizes.

        Returns a list of dicts, one dict per image. The dict contains:
        rois: [N, (y1, x1, y2, x2)] detection bounding boxes
        class_ids: [N] int class IDs
        scores: [N] float probability scores for the class IDs
        masks: [H, W, N] instance binary masks
        """
        assert self.mode == "inference", "Create model in inference mode."
        assert len(images) == self.config.BATCH_SIZE, "len(images) must be equal to BATCH_SIZE"

        # Mold inputs to format expected by the neural network
        molded_images, image_metas, windows = self.mold_inputs(images)

        # Validate image sizes
        # All images in a batch MUST be of the same size
        image_shape = molded_images[0].shape
        for g in molded_images[1:]:
            assert g.shape == image_shape, \
                "After resizing, all images must have the same size. Check IMAGE_RESIZE_MODE and image sizes."

        # Anchors
        anchors = self.get_anchors(image_shape)
        # Duplicate across the batch dimension because Keras requires it
        # TODO: can this be optimized to avoid duplicating the anchors?
        anchors = np.broadcast_to(anchors, (self.config.BATCH_SIZE,) + anchors.shape)

        # Run object detection
        detections, mrcnn_class_logits, mrcnn_class, mrcnn_bbox, mrcnn_mask, rpn_rois, rpn_class, rpn_bbox = \
            self.sess.run(self.model_output,
                          feed_dict={self.model_input[0]: molded_images,
                                     self.model_input[1]: image_metas,
                                     self.model_input[2]: anchors})
        # Process detections
        results = []
        for i, image in enumerate(images):
            final_rois, final_class_ids, final_scores, final_masks = \
                self.unmold_detections(detections[i], mrcnn_mask[i],
                                       image.shape, molded_images[i].shape,
                                       windows[i])
            results.append({
                "rois": final_rois,
                "class_ids": final_class_ids,
                "scores": final_scores,
                "masks": final_masks,
            })
        return results

    def detect1(self, train_dataset):
        data = data_generator(train_dataset, self.config, shuffle=True,
                              batch_size=1)

        single_input, _ = data.__next__()

        molded_images = single_input[0]
        image_metas = single_input[1]
        windows = [image_metas[0][7:11].astype(np.int32)]
        original_shape = image_metas[0][1:4].astype(np.int32)

        # Validate image sizes
        # All images in a batch MUST be of the same size
        image_shape = molded_images[0].shape
        for g in molded_images[1:]:
            assert g.shape == image_shape, \
                "After resizing, all images must have the same size. Check IMAGE_RESIZE_MODE and image sizes."

        # Anchors
        anchors = self.get_anchors(image_shape)
        # Duplicate across the batch dimension because Keras requires it
        # TODO: can this be optimized to avoid duplicating the anchors?
        anchors = np.broadcast_to(anchors, (self.config.BATCH_SIZE,) + anchors.shape)

        # Run object detection
        detections, mrcnn_class_logits, mrcnn_class, mrcnn_bbox, mrcnn_mask, rpn_rois, rpn_class, rpn_bbox = \
            self.sess.run(self.model_output,
                          feed_dict={self.model_input[0]: molded_images,
                                     self.model_input[1]: image_metas,
                                     self.model_input[2]: anchors})
        # dic = {}
        # dic['rpn_roi_200'] = rpn_rois
        # np.save('rpn_roi_detect_dic200.npy', [dic])
        # Process detections
        results = []
        i = 0
        final_rois, final_class_ids, final_scores, final_masks = \
            self.unmold_detections(detections[i], mrcnn_mask[i],
                                   original_shape, molded_images[i].shape,
                                   windows[i])
        print(final_class_ids)
        original_image = unmold_image(molded_images[0], self.config)[windows[0][0]:windows[0][2], windows[0][1]:windows[0][3], :]
        original_image = (resize(original_image, output_shape=original_shape) * 255).astype(np.uint8)
        results.append({
            "images": original_image,
            "rois": final_rois,
            "class_ids": final_class_ids,
            "scores": final_scores,
            "masks": final_masks,
        })
        return results

    def detect_load_weight(self, path):
        # saver = tf.train.Saver()
        self.saver.restore(self.sess, path)

