Using segmentation masks (png files) to train detectron2

[siwalex]

Hello,

I have the same problem than this guys : #742

i use this colab (official) : https://colab.research.google.com/drive/16jcaJoc6bCFAQ96jDe2HwtXj7BMD_-m5

and i want to train a dataset from masks. not labeling.

i'm sorry but i don't understand the issue200 : #200

that is suppose to resolve my problem.

Can you explain to me how i'm suppose to do/change in the colab to train with masks?

Thanks.
Alexandre

[abdulazizm]

@siwalex : I am also trying to do the same using binary mask (from A2D2 instance segmentation annotations) directly. Tried converting binary mask to RLE but seems to be a tidy process since I have more instances from different classes in a single binary mask. Let me know if you find any workarounds

[siwalex]

hello, sorry but i didn't find any issue for my problem
i made my own mask to json script

[abdulazizm]

how did you split each instance from bitmask (with multiple instances from multiple classes)?

[AKMourato]

@abdulazizm here goes a simple script to generate labels in coco format from pgn binary segmentation masks (multiple instances) for instance segmentation.
Just make sure the structure goes as follows:
e.g:

└───train
    │ 
    └───annotations
    │    │ <image_id>_<object_class_name>_<annotation_id>.png
    │    │ ...
    │   
    └───images
         │   <image_id>.jpeg
         │   ...
└───val
(....)
└───test

import datetime
import json
import os
import re
import fnmatch
from PIL import Image
import numpy as np
from pycococreator.pycococreatortools import pycococreatortools

ROOT_DIR = '/train/'
IMAGE_DIR = 
ANNOTATION_DIR = 

INFO = {
    "description": "datasetX",
    "url": "",
    "version": "1.0",
    "year": 2021,
    "contributor": "AKM22",
    "date_created": datetime.datetime.utcnow().isoformat(' ')
}

LICENSES = [
    {
        "id": 1,
        "name": "Attribution-NonCommercial-ShareAlike License",
        "url": "http://creativecommons.org/licenses/by-nc-sa/2.0/"
    }
]

CATEGORIES = [
    {
        'id': 1,
        'name': 'XYZ',
        'supercategory': 'omega',
    }]
def filter_for_jpeg(root, files):
    file_types = ['*.png']
    file_types = r'|'.join([fnmatch.translate(x) for x in file_types])
    files = [os.path.join(root, f) for f in files]
    files = [f for f in files if re.match(file_types, f)]
    
    return files

def filter_for_annotations(root, files, image_filename):
    file_types = ['*.png']
    file_types = r'|'.join([fnmatch.translate(x) for x in file_types])
    basename_no_extension = os.path.splitext(os.path.basename(image_filename))[0]
    file_name_prefix = basename_no_extension + '.*'
    files = [os.path.join(root, f) for f in files]
    files = [f for f in files if re.match(file_types, f)]
    files = [f for f in files if re.match(file_name_prefix, os.path.splitext(os.path.basename(f))[0])]

    return files

def main():

    coco_output = {
        "info": INFO,
        "licenses": LICENSES,
        "categories": CATEGORIES,
        "images": [],
        "annotations": []
    }

    image_id = 1
    segmentation_id = 1
    
    # filter for jpeg images
    for root, _, files in os.walk(IMAGE_DIR):
        image_files = filter_for_jpeg(root, files)

        # go through each image
        for image_filename in sorted(image_files):
            image = Image.open(image_filename)
            image_info = pycococreatortools.create_image_info(
                image_id, os.path.basename(image_filename), image.size)
            coco_output["images"].append(image_info)

            # filter for associated png annotations
            for root, _, files in os.walk(ANNOTATION_DIR):
                annotation_files = filter_for_annotations(root, files, image_filename)

                # go through each associated annotation
                for annotation_filename in sorted(annotation_files):
                    
                    print(annotation_filename)
                    class_id = [x['id'] for x in CATEGORIES if x['name'] in annotation_filename][0]

                    category_info = {'id': class_id, 'is_crowd': 'crowd' in image_filename}
                    binary_mask = np.asarray(Image.open(annotation_filename)
                        .convert('1')).astype(np.uint8)
                    
                    annotation_info = pycococreatortools.create_annotation_info(
                        segmentation_id, image_id, category_info, binary_mask,
                        image.size, tolerance=2)

                    if annotation_info is not None:
                        coco_output["annotations"].append(annotation_info)

                    segmentation_id = segmentation_id + 1

            image_id = image_id + 1

    with open('{}/instances_images.json'.format(ROOT_DIR), 'w') as output_json_file:
        json.dump(coco_output, output_json_file)


if __name__ == "__main__":
    main()

[abdulazizm]

@AKMourato Thanks for the complete code and your engagement. I modified it a little to create panoptic JSON files too (in addition to instance JSON file). Hope I can proceed with registering the dataset for training panoptic deeplab (as mentioned here -> #1691 (comment) ) and proceed further. Let me know your suggestions or feedback if any.

[jiunyen-ching]

Training on detectron2 for instance segmentation. Encoding of bitmasks is using RLE instead of polygons.

I read in some articles that when encoding instance masks (that has holes) via polygon annotations, they need to be treated with extra steps to account for the hole contours. I am not unsure that I won't break anything while modifying detectron's backend codes to achieve this so I chose the encode the masks using RLE. The masks can be encoded and decoded properly while preserving the holes.

Modifying this section from the balloon example from detectron's colab notebook with the dataset structure like this:

└─── your_dataset
      └─── train 
           └─── [image_id]_color.jpg
           └─── [image_id]_instances.png
           └─── [image_id]_class_labels.png
      └─── val
           └─── [image_id]_color.jpg
           └─── [image_id]_instances.png
           └─── [image_id]_class_labels.png

import pycocotools
from detectron2.structures import BoxMode

# img_dir should be '/path/to/your_dataset/train' or '/path/to/your_dataset/val'
def get_nyu_dicts(img_dir):

    # example of instance filename: 000042_instances.png
    filelist = [x[:6] for x in sorted(os.listdir(img_dir)) if x.endswith('_instances.png')]

    dataset_dicts = []
    for idx, file in enumerate(filelist):
        record = {}

        # Use RGB images
        filename = os.path.join(img_dir, file + '_color.jpg')
        height, width = cv2.imread(filename).shape[:2]

        record["file_name"] = filename
        record["image_id"] = idx
        record["height"] = height
        record["width"] = width

        # annos = v["regions"] # replaced by instance ids + class labels
        objs = []

        # Load files
        instance_img = cv2.imread(os.path.join(img_dir, file + '_instances.png'), -1)
        label_img    = cv2.imread(os.path.join(img_dir, file + '_class_labels.png'), -1)

        # Get instances
        instance_ids = np.unique(instance_img)
        if instance_ids[0] == 0: # avoid unlabelled areas
            instance_ids = instance_ids[1:]

        # Pair instance with label
        for _id in instance_ids:
            mask = instance_img == _id
            py, px = np.where(mask == True) # bounding box
            label = label_img[mask][0] # [0] to select element from array         
            mask = pycocotools.mask.encode(mask.astype(np.uint8, order="F"))
    
            obj = {
                "bbox": [np.min(px), np.min(py), np.max(px), np.max(py)],
                "bbox_mode": BoxMode.XYXY_ABS,
                "segmentation": mask,
                "category_id": label-1, # my dataset starts with 1 as a class label so I do "label-1" to let category id start from 0
            }
            objs.append(obj)
        record["annotations"] = objs
        dataset_dicts.append(record)
    return dataset_dicts

Followed up by registering dataset

path = '/path/to/your_dataset'

for d in ["train", "test"]:
    DatasetCatalog.register("nyu_" + d, lambda d=d: get_nyu_dicts(path + "/" + d))
    MetadataCatalog.get("nyu_" + d).set(thing_classes=["ceiling", "floor", "wall", "window", "chair", "bed", "sofa", "table", "tvs", "furniture", "objects"])
nyu_metadata = MetadataCatalog.get("nyu_train")

Also add this line to use bitmask
cfg.INPUT.MASK_FORMAT='bitmask'

Overall config should look something like this

from detectron2.engine import DefaultTrainer

cfg = get_cfg()
cfg.INPUT.MASK_FORMAT='bitmask' # for rle
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("nyu_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
cfg.SOLVER.MAX_ITER = 300    # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
cfg.SOLVER.STEPS = []        # do not decay learning rate
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 11  # only has one class (balloon). (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)
# NOTE: this config means the number of classes, but a few popular unofficial tutorials incorrect uses num_classes+1 here.

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

[norton-chris]

If you are using binary masks, I used this repo to convert my binary segmentation masks into coco json files:
https://github.com/brunobelloni/binary-to-coco-json-converter