text: """
  Mask R-CNN
  """
  if __name__ == '__main__':
      import matplotlib
      # Agg backend runs without a display
      matplotlib.use('Agg')
      import matplotlib.pyplot as plt
  
  import datetime
  import os
  import sys
  
  import numpy as np
  import scipy.io
  from imgaug import augmenters as iaa
  
  # Root directory of the project
  ROOT_DIR = os.path.abspath(".")
  
  print(ROOT_DIR)
  # Import Mask RCNN
  sys.path.append(ROOT_DIR)  # To find local version of the library
  from mrcnn.config import Config
  from mrcnn import model as modellib, utils
  from mrcnn import visualize
  
  os.environ['CUDA_VISIBLE_DEVICES'] = '0, 2'
  # Path to trained weights file
  COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
  
  # Directory to save logs and model checkpoints, if not provided
  # through the command line argument --logs
  DEFAULT_LOGS_DIR = os.path.join(ROOT_DIR, "logs")
  RESULTS_DIR = os.path.join(ROOT_DIR, "results/consep/")
  
  CONSEP_CLASS_IDS = [
      'other',
      'inflammatory',
      'healthy_epithelial',
      'dysplastic/malignant_epithelial',
      'fibroblast',
      'muscle',
      'endothelial'
  ]
  
  
  ############################################################
  #  Configurations
  ############################################################
  
  
  class CoNSepConfig(Config):
      """Configuration for training on CoNSep.
      Derives from the base Config class and overrides values specific
      to the CoNSep dataset.
      """
      # Give the configuration a recognizable name
      NAME = "consep"
  
      # Adjust down if you use a smaller GPU depending on your GPU memory.
      IMAGES_PER_GPU = 2  # 6
  
      # Uncomment to train on 8 GPUs (default is 1)
      GPU_COUNT = 2
  
      # Number of classes (including background)
      NUM_CLASSES = 1 + 7  # Background + nucleus
  
      # Number of training and validation steps per epoch
      STEPS_PER_EPOCH = 27 // IMAGES_PER_GPU
      VALIDATION_STEPS = max(1, 14 // IMAGES_PER_GPU)
  
      # Don't exclude based on confidence. Since we have two classes. 0 for nucleus
      # then 0.5 is the minimum anyway as it picks between nucleus and BG
      DETECTION_MIN_CONFIDENCE = 0.4  # ROIs
  
      # Backbone network architecture
      # Supported values are: resnet50, resnet101
      BACKBONE = "resnet50"
  
      # Input image resizing
      # Random crops of size 256x256
      IMAGE_RESIZE_MODE = "crop"
      IMAGE_MIN_DIM = 256
      IMAGE_MAX_DIM = 256
      IMAGE_MIN_SCALE = 0  # 2.0
  
      # Length of square anchor side in pixels
      RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)
  
      # ROIs kept after non-maximum supression (training and inference)
      POST_NMS_ROIS_TRAINING = 1000
      POST_NMS_ROIS_INFERENCE = 2000
  
      # Non-max suppression threshold to filter RPN proposals.
      # You can increase this during training to generate more propsals.
      RPN_NMS_THRESHOLD = 0.8
  
      # How many anchors per image to use for RPN training
      RPN_TRAIN_ANCHORS_PER_IMAGE = 64
  
      # Image mean (RGB)
      MEAN_PIXEL = np.array([43.53, 39.56, 48.22])
  
      # If enabled, resizes instance masks to a smaller size to reduce
      # memory load. Recommended when using high-resolution images.
      USE_MINI_MASK = True
      MINI_MASK_SHAPE = (56, 56)  # (height, width) of the mini-mask
  
      # Number of ROIs per image to feed to classifier/mask heads
      # The Mask RCNN paper uses 512 but often the RPN doesn't generate
      # enough positive proposals to fill this and keep a positive:negative
      # ratio of 1:3. You can increase the number of proposals by adjusting
      # the RPN NMS threshold.
      TRAIN_ROIS_PER_IMAGE = 128
  
      # Percent of positive ROIs used to train classifier/mask heads
      ROI_POSITIVE_RATIO = 0.4
  
      # Maximum number of ground truth instances to use in one image
      MAX_GT_INSTANCES = 200
  
      # Max number of final detections per image
      DETECTION_MAX_INSTANCES = 400
  
  
  ############################################################
  #  Dataset
  ############################################################
  
  class CoNSepDataset(utils.Dataset):
      def load_consep(self, dataset_dir, subset, return_consep=False):
          """Load a subset of the CoNSep dataset.
          dataset_dir: The root directory of the CoNSep dataset.
          subset: What to load (train, test)
          return_coco: If True, returns the COCO object.
          auto_download: Automatically download and unzip MS-COCO images and annotations
          """
  
          image_dir = "{}/Images/".format(subset)
          dataset_dir = os.path.join(dataset_dir, image_dir)
          # print(dataset_dir, os.walk(dataset_dir))
          # Get image ids from directory names
          image_ids = os.listdir(dataset_dir)#next(os.walk(dataset_dir))[-1]
          print(image_ids)
          # Add classes
          for i, name in enumerate(CONSEP_CLASS_IDS):
              self.add_class("consep", i + 1, name)
  
          # Add images
          for image_id in image_ids:
              self.add_image(
                  "consep",
                  image_id=image_id,
                  path=os.path.join(dataset_dir, image_id)
              )
          # print(self.image_info)
  
      def load_mask(self, image_id):
          """Generate instance masks for an image.
          Returns:
          masks: A bool array of shape [height, width, instance count] with
              one mask per instance.
          class_ids: a 1D array of class IDs of the instance masks.
          """
          # print(self.image_info, image_id)
          info = self.image_info[image_id]
          # Get mask directory from image path
          mask_dir = os.path.join(os.path.dirname(os.path.dirname(info['path'])), "Labels/{}.mat".format(info['id'].replace('.png', '')))
          # Read mask files from .mat image
          '''
          'inst_map' is a 1000x1000 array containing a unique integer for each individual nucleus. i.e the map ranges from 0 to N, where 0 is the background and N is the number of nuclei.
          'type_map' is a 1000x1000 array where each pixel value denotes the class of that pixel. The map ranges from 0 to 7, where 7 is the total number of classes in CoNSeP.       
          'inst_type' is a Nx1 array, indicating the type of each instance (in order of inst_map ID)       
          'inst_centroid' is a Nx2 array, giving the x and y coordinates of the centroids of each instance (in order of inst map ID).
          '''
          mask = []
          m = scipy.io.loadmat(mask_dir)
          inst_map = m['inst_map']
          class_ids = m['inst_type']
          num_inst = len(class_ids)
          for i in range(num_inst):
              new = np.where(inst_map == i + 1, 1, 0).astype(np.bool)
              mask.append(new)
          mask = np.stack(mask, axis=-1)
          class_ids = np.array(np.squeeze(class_ids), dtype=np.int32)
          return mask, class_ids
  
      def image_reference(self, image_id):
          """Return the path of the image."""
          info = self.image_info[image_id]
          if info["source"] == "consep":
              return info["id"]
          else:
              super(self.__class__, self).image_reference(image_id)
  
  
  ############################################################
  #  CoNSep Evaluation
  ############################################################
  
  def rle_encode(mask):
      """Encodes a mask in Run Length Encoding (RLE).
      Returns a string of space-separated values.
      """
      assert mask.ndim == 2, "Mask must be of shape [Height, Width]"
      # Flatten it column wise
      m = mask.T.flatten()
      # Compute gradient. Equals 1 or -1 at transition points
      g = np.diff(np.concatenate([[0], m, [0]]), n=1)
      # 1-based indicies of transition points (where gradient != 0)
      rle = np.where(g != 0)[0].reshape([-1, 2]) + 1
      # Convert second index in each pair to lenth
      rle[:, 1] = rle[:, 1] - rle[:, 0]
      return " ".join(map(str, rle.flatten()))
  
  
  def rle_decode(rle, shape):
      """Decodes an RLE encoded list of space separated
      numbers and returns a binary mask."""
      rle = list(map(int, rle.split()))
      rle = np.array(rle, dtype=np.int32).reshape([-1, 2])
      rle[:, 1] += rle[:, 0]
      rle -= 1
      mask = np.zeros([shape[0] * shape[1]], np.bool)
      for s, e in rle:
          assert 0 <= s < mask.shape[0]
          assert 1 <= e <= mask.shape[0], "shape: {}  s {}  e {}".format(shape, s, e)
          mask[s:e] = 1
      # Reshape and transpose
      mask = mask.reshape([shape[1], shape[0]]).T
      return mask
  
  
  def mask_to_rle(image_id, mask, scores):
      "Encodes instance masks to submission format."
      assert mask.ndim == 3, "Mask must be [H, W, count]"
      # If mask is empty, return line with image ID only
      if mask.shape[-1] == 0:
          return "{},".format(image_id)
      # Remove mask overlaps
      # Multiply each instance mask by its score order
      # then take the maximum across the last dimension
      order = np.argsort(scores)[::-1] + 1  # 1-based descending
      mask = np.max(mask * np.reshape(order, [1, 1, -1]), -1)
      # Loop over instance masks
      lines = []
      for o in order:
          m = np.where(mask == o, 1, 0)
          # Skip if empty
          if m.sum() == 0.0:
              continue
          rle = rle_encode(m)
          lines.append("{}, {}".format(image_id, rle))
      return "\n".join(lines)
  
  
  ############################################################
  #  Training
  ############################################################
  
  def train(model, dataset_dir, subset):
      """Train the model."""
      # Training dataset.
      dataset_train = CoNSepDataset()
      dataset_train.load_consep(dataset_dir, subset)
      dataset_train.prepare()
  
      # Validation dataset
      dataset_val = CoNSepDataset()
      dataset_val.load_consep(dataset_dir, "test")
      dataset_val.prepare()
  
      # Image augmentation
      # http://imgaug.readthedocs.io/en/latest/source/augmenters.html
      augmentation = iaa.SomeOf((0, 2), [
          iaa.Fliplr(0.5),
          iaa.Flipud(0.5),
          iaa.OneOf([iaa.Affine(rotate=90),
                     iaa.Affine(rotate=180),
                     iaa.Affine(rotate=270)]),
          iaa.Multiply((0.8, 1.5)),
          iaa.GaussianBlur(sigma=(0.0, 5.0))
      ])
  
      # *** This training schedule is an example. Update to your needs ***
  
      # If starting from imagenet, train heads only for a bit
      # since they have random weights
      print("Train network heads")
      model.train(dataset_train, dataset_val,
                  learning_rate=config.LEARNING_RATE,
                  epochs=20,
                  augmentation=augmentation,
                  layers='heads')
  
      print("Train all layers")
      model.train(dataset_train, dataset_val,
                  learning_rate=config.LEARNING_RATE,
                  epochs=40,
                  augmentation=augmentation,
                  layers='all')
  
  
  ############################################################
  #  Detection
  ############################################################
  
  def evaluate(model, dataset_dir, subset):
      """Run detection on images in the given directory."""
      print("Running on {}".format(dataset_dir))
  
      # Create directory
      if not os.path.exists(RESULTS_DIR):
          os.makedirs(RESULTS_DIR)
      submit_dir = "submit_{:%Y%m%dT%H%M%S}".format(datetime.datetime.now())
      submit_dir = os.path.join(RESULTS_DIR, submit_dir)
      os.makedirs(submit_dir)
  
      # Read dataset
      dataset = CoNSepDataset()
      dataset.load_consep(dataset_dir, subset)
      dataset.prepare()
      # Load over images
      submission = []
      for image_id in dataset.image_ids:
          # Load image and run detection
          image = dataset.load_image(image_id)
          # Detect objects
          r = model.detect([image], verbose=0)[0]
          # Encode image to RLE. Returns a string of multiple lines
          source_id = dataset.image_info[image_id]["id"]
          rle = mask_to_rle(source_id, r["masks"], r["scores"])
          submission.append(rle)
          # Save image with masks
          visualize.display_instances(
              image, r['rois'], r['masks'], r['class_ids'],
              dataset.class_names, r['scores'],
              show_bbox=False, show_mask=False,
              title="Predictions")
          plt.savefig("{}/{}.png".format(submit_dir, dataset.image_info[image_id]["id"]))
  
      # Save to csv file
      submission = "ImageId,EncodedPixels\n" + "\n".join(submission)
      file_path = os.path.join(submit_dir, "submit.csv")
      with open(file_path, "w") as f:
          f.write(submission)
      print("Saved to ", submit_dir)
  
  
  if __name__ == '__main__':
      import argparse
  
      # Parse command line arguments
      parser = argparse.ArgumentParser(
          description='Train Mask R-CNN on CoNSep.')
      parser.add_argument("command",
                          metavar="<command>",
                          help="'train' or 'evaluate' on MS COCO")
      parser.add_argument('--dataset', required=True,
                          metavar="/path/to/dataset/",
                          help='Directory of the CoNSep dataset')
      parser.add_argument('--weights', required=True,
                          metavar="/path/to/weights.h5",
                          help="Path to weights .h5 file or 'coco'")
      parser.add_argument('--logs', required=False,
                          default=DEFAULT_LOGS_DIR,
                          metavar="/path/to/logs/",
                          help='Logs and checkpoints directory (default=logs/)')
      parser.add_argument('--subset', required=False,
                          metavar="Dataset sub-directory",
                          help="Subset of dataset to run prediction on")
      parser.add_argument('--limit', required=False,
                          default=500,
                          metavar="<image count>",
                          help='Images to use for evaluation (default=500)')
  
      args = parser.parse_args()
      # Validate arguments
      if args.command == "train":
          assert args.dataset, "Argument --dataset is required for training"
      elif args.command == "evaluate":
          assert args.subset, "Provide --subset to run prediction on"
  
      print("Weights: ", args.weights)
      print("Dataset: ", args.dataset)
      if args.subset:
          print("Subset: ", args.subset)
      print("Logs: ", args.logs)
  
      # Configurations
      if args.command == "train":
          config = CoNSepConfig()
      else:
          class InferenceConfig(CoNSepConfig):
              # Set batch size to 1 since we'll be running inference on
              # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
              GPU_COUNT = 1
              IMAGES_PER_GPU = 1
              DETECTION_MIN_CONFIDENCE = 0
              # Don't resize imager for inferencing ??
              IMAGE_RESIZE_MODE = "pad64"
  
  
          config = InferenceConfig()
      config.display()
  
      # Create model
      if args.command == "train":
          model = modellib.MaskRCNN(mode="training", config=config,
                                    model_dir=args.logs)
      else:
          model = modellib.MaskRCNN(mode="inference", config=config,
                                    model_dir=args.logs)
  
      # Select weights file to load
      if args.weights.lower() == "coco":
          weights_path = COCO_MODEL_PATH
          # Download weights file
          if not os.path.exists(weights_path):
              utils.download_trained_weights(weights_path)
      elif args.weights.lower() == "last":
          # Find last trained weights
          weights_path = model.find_last()
      elif args.weights.lower() == "imagenet":
          # Start from ImageNet trained weights
          weights_path = model.get_imagenet_weights()
      else:
          weights_path = args.weights
  
      # Load weights
      print("Loading weights ", weights_path)
      if args.weights.lower() == "coco":
          # Exclude the last layers because they require a matching
          # number of classes
          model.load_weights(weights_path, by_name=True, exclude=[
              "mrcnn_class_logits", "mrcnn_bbox_fc",
              "mrcnn_bbox", "mrcnn_mask"])
      else:
          model.load_weights(weights_path, by_name=True)
  
      # Train or evaluate
      if args.command == "train":
          train(model, args.dataset, args.subset)
      elif args.command == "evaluate":
          evaluate(model, args.dataset, args.subset)
      else:
          print("'{}' is not recognized. "
                "Use 'train' or 'evaluate'".format(args.command))
  