Maskrcnn-benchmark is an open-source Facebook Benchmark algorithm project that includes algorithms for detection, segmentation, and human keypoints.

This series consists of two parts:

  • Chapter 1 Setting up the environment;
  • Part II Training and Validation;

training

Using maskrCNN-benchmark training model, you can refer to.

Data set:

  • Annotations download complete COCO dataset: Annotations, Test2014, Train2014, VAL2014;
  • Download the COCO minivalidation sets provided by FAIR: Minival and ValminusMinival;

Select the training template: e2e_MASk_RCNn_R_50_fpn_1x. yaml, where:

WEIGHT: "catalog://ImageNetPretrained/MSRA/R-50"  # Pre-training weights
DATASETS:  # data set
  TRAIN: ("coco_2014_train"."coco_2014_valminusminival")
  TEST: ("coco_2014_minival",)
MAX_ITER: 90000  # Maximum training rounds
Copy the code

Other parameters setting location: maskrcnn_benchmark/config/defaults. Py

Such as:

  • _C.SOLVER.CHECKPOINT_PERIOD = 2500Save the rounds;
  • _C.SOLVER.IMS_PER_BATCH = 16, trainingbatch_size;
  • _C.OUTPUT_DIR = "./models", model output path;

Specifies the number of gpus:

export NGPUS=4
Copy the code

Training model:

python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_net.py --config-file "configs/e2e_mask_rcnn_R_50_FPN_1x.yaml"

nohup python -u -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_net.py --config-file "configs/e2e_mask_rcnn_R_50_FPN_1x.yaml" &
Copy the code

The output model is in./ Models, and the last model is model_0090,000.pth.

test

Change the WEIGHT of configuration file e2e_mask_rCNn_R_50_fpn_1x.my. yaml to the model path, for example, WEIGHT: “model_0090000.pth”.

The test logic is as follows:

def main():
    img_path = os.path.join(DATA_DIR, 'aoa-mina.jpeg')

    img = cv2.imread(img_path)
    print('[Info] img size: {}'.format(img.shape))
    
    config_file = ".. /configs/e2e_mask_rcnn_R_50_FPN_1x.my.yaml"

    cfg.merge_from_file(config_file)  Set the configuration file
    cfg.merge_from_list(["MODEL.MASK_ON", True])
    cfg.merge_from_list(["MODEL.DEVICE"."cpu"])  # specify CPU

    coco_demo = COCODemo(  Create a model file
        cfg,
        # show_mask_heatmaps=True,Min_image_size = 800, confidence_threshold = 0.7, ) predictions = coco_demo.compute_prediction(img) top_predictions = coco_demo.select_top_predictions(predictions) show_mask(img, top_predictions, coco_demo)print('Execution complete! ')
Copy the code

Show_mask shows the effect of segmentation and draws polygon boxes based on matplotlib’s image:

def show_mask(img, predictions, coco_demo):
    """ Show Partitioning effects: Param img: Numpy Images: Param Predictions: Partitioning Results: Param Coco_demo: Function Set :return: Show Partitioning effects ""
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    pylab.rcParams['figure.figsize'] = (8.0.10.0)  # Image size
    plt.imshow(img)  The image needs to be filled in advance
    # plt.show()

    extra_fields = predictions.extra_fields
    masks = extra_fields['mask']
    labels = extra_fields['labels']
    name_list = [coco_demo.CATEGORIES[l] for l in labels]

    seg_list = []
    for mask in masks:
        mask = torch.squeeze(mask)
        segmentation = binary_mask_to_polygon(mask, tolerance=2) [0]
        seg_list.append(segmentation)

    ax = plt.gca()
    ax.set_autoscale_on(False)

    polygons, color = [], []

    np.random.seed(37)

    for name, seg in zip(name_list, seg_list):
        c = (np.random.random((1.3)) * 0.8 + 0.2).tolist()[0]

        poly = np.array(seg).reshape((int(len(seg) / 2), 2))
        c_x, c_y = get_center_of_polygon(poly)  # Calculate the center of the polygon

        # 0 to 26 are large categories, the rest are small categories and each label is drawn only once
        tc = c - np.array([0.5.0.5.0.5])  # drop color
        tc = np.maximum(tc, 0.0)  The minimum value is 0

        plt.text(c_x, c_y, name, ha='left', wrap=True, color=tc,
                 bbox=dict(facecolor='white', alpha=0.5))  # draw the tag

        polygons.append(pylab.Polygon(poly))  # draw polygons
        color.append(c)

    p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.4)  # add polygons
    ax.add_collection(p)
    p = PatchCollection(polygons, facecolor='none', edgecolors=color, linewidths=2)  # Add polygon box
    ax.add_collection(p)

    plt.axis('off')
    ax.get_xaxis().set_visible(False)  # this removes the ticks and numbers for x axis
    ax.get_yaxis().set_visible(False)  # this removes the ticks and numbers for y axis

    out_folder = os.path.join(ROOT_DIR, 'demo'.'out')
    mkdir_if_not_exist(out_folder)
    out_file = os.path.join(out_folder, 'test.png'.format())
    plt.savefig(out_file, bbox_inches='tight', pad_inches=0, dpi=200)

    plt.close()  # Avoid drawing all images together
Copy the code

The effect is as follows:

OK, that’s all!