text: import os 
  import sys
  import numpy as np 
  import skimage.draw
  import matplotlib.pyplot as plt
  
  import warnings
  warnings.filterwarnings("ignore")
  
  # Root directory of the project
  ROOT_DIR = os.path.abspath("../../")
  
  # Import Mask RCNN
  sys.path.append(ROOT_DIR) 
  
  from mrcnn import model as modellib, utils, visualize
  from pcb import pcbDataset, pcbConfig
  
  
  # Path to trained weights file
  COCO_WEIGHTS_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")
  
  
  def get_ax(rows=1, cols=1, size=16):
      """Return a Matplotlib Axes array to be used in
      all visualizations in the notebook. Provide a 
      central point to control graph sizes.
  
      Adjust the size attribute to control how big to render images
      """
      _, ax = plt.subplots(rows, cols, figsize=(size*cols, size*rows))
      return ax
  
  
  def train(model):
      print(args)
      """Train the model."""
      # Training dataset.
      dataset_train = pcbDataset()
      dataset_train.load_pcb(args.dataset, "train")
      dataset_train.prepare()
  
      # Validation dataset
      dataset_val = pcbDataset()
      dataset_val.load_pcb(args.dataset, "val")
      dataset_val.prepare()
  
      # *** This training schedule is an example. Update to your needs ***
      # Since we're using a very small dataset, and starting from
      # COCO trained weights, we don't need to train too long. Also,
      # no need to train all layers, just the heads should do it.
      print("Training network heads")
      model.train(dataset_train, dataset_val,
                  learning_rate=config.LEARNING_RATE,
                  epochs=50,
                  layers='heads')
  
      # Training - Stage 2
      # Finetune layers from ResNet stage 4 and up
      print("Fine tune Resnet stage 4 and up")
      model.train(dataset_train, dataset_val,
                  learning_rate=config.LEARNING_RATE,
                  epochs=120,
                  layers='4+')
  
      # Training - Stage 3
      # Fine tune all layers
      print("Fine tune all layers")
      model.train(dataset_train, dataset_val,
                  learning_rate=config.LEARNING_RATE / 10,
                  epochs=160,
                  layers='all')
  
  
  def test(model, image_path):
      print("Running on {}".format(image_path))
      
      # Read image
      image = skimage.io.imread(image_path)
  
      # detect objects
      r = model.detect([image], verbose=1)[0]
  
      # Display results
      ax = get_ax(1)
      visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
                                  ["Background", "IC"], r['scores'], ax=ax,
                                  title="Predictions")
  
  
  ############################################################
  #  Training
  ############################################################
  
  if __name__ == '__main__':
      import argparse
  
      # Parse command line arguments
      parser = argparse.ArgumentParser(
          description='Train Mask R-CNN to detect ICs.')
      parser.add_argument("command",
                          metavar="<command>",
                          help="'train' or 'test'")
      parser.add_argument('--dataset', required=False,
                          metavar="/path/to/pcb/dataset/",
                          help='Directory of the pcb 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('--image', required=False,
                          metavar="path or URL to image",
                          help='Image to apply the detector on')
      args = parser.parse_args()
  
      # Validate arguments
      if args.command == "train":
          assert args.dataset, "Argument --dataset is required for training"
      elif args.command == "test":
          assert args.image,\
                 "Provide --image or --video to apply color splash"
  
      print("Weights: ", args.weights)
      print("Dataset: ", args.dataset)
      print("Logs: ", args.logs)
  
      # Configurations
      if args.command == "train":
          config = pcbConfig()
      else:
          class InferenceConfig(pcbConfig):
              # 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
          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_WEIGHTS_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()[1]
      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)
      elif args.command == "test":
          test(model, image_path=args.image)
      else:
          print("'{}' is not recognized. "
                "Use 'train' or 'test'".format(args.command))