dataset - 如何从黑白蒙版生成 COCO 数据集

问题描述

我有一个由焊缝和掩码组成的数据集（焊缝为白色，背景为黑色），尽管我需要使用 Mask R-CNN，所以我必须将它们转换为 COCO 数据集注释。有人对如何做到这一点有任何建议吗？

我试过这个：https ://github.com/chrise96/image-to-coco-json-converter

但我收到此错误：

---------------------------------------------------------------------------

KeyError                                  Traceback (most recent call last)

<ipython-input-3-0ddc235b1528> in <module>
     94 
     95         # Create images and annotations sections
---> 96         coco_format["images"], coco_format["annotations"], annotation_cnt = images_annotations_info(mask_path)
     97 
     98         with open("output/{}.json".format(keyword),"w") as outfile:

<ipython-input-3-0ddc235b1528> in images_annotations_info(maskpath)
     57         sub_masks = create_sub_masks(mask_image_open, w, h)
     58         for color, sub_mask in sub_masks.items():
---> 59             category_id = category_colors[color]
     60 
     61             # "annotations" info

KeyError: '(1, 1, 1)'

这是代码，我刚刚添加了焊接分类：


import glob

from src.create_annotations import *

# Label ids of the dataset
category_ids = {
    "outlier": 0,
    "window": 1,
    "wall": 2,
    "balcony": 3,
    "door": 4,
    "roof": 5,
    "sky": 6,
    "shop": 7,
    "chimney": 8,
    "weld": 9,
}

# Define which colors match which categories in the images
category_colors = {
    "(0, 0, 0)": 0, # Outlier
    "(255, 0, 0)": 1, # Window
    "(255, 255, 0)": 2, # Wall
    "(128, 0, 255)": 3, # Balcony
    "(255, 128, 0)": 4, # Door
    "(0, 0, 255)": 5, # Roof
    "(128, 255, 255)": 6, # Sky
    "(0, 255, 0)": 7, # Shop
    "(128, 128, 128)": 8, # Chimney
    "(255, 255, 255)": 9 # Weld
}

# Define the ids that are a multiplolygon. In our case: wall, roof and sky
multipolygon_ids = [9, 2, 5, 6]

# Get "images" and "annotations" info
def images_annotations_info(maskpath):
    # This id will be automatically increased as we go
    annotation_id = 0
    image_id = 0
    annotations = []
    images = []

    for mask_image in glob.glob(maskpath + "*.png"):
        # The mask image is *.png but the original image is *.jpg.
        # We make a reference to the original file in the COCO JSON file
        original_file_name = os.path.basename(mask_image).split(".")[0] + ".jpg"

        # Open the image and (to be sure) we convert it to RGB
        mask_image_open = Image.open(mask_image).convert("RGB")
        w, h = mask_image_open.size

        # "images" info
        image = create_image_annotation(original_file_name, w, h, image_id)
        images.append(image)

        sub_masks = create_sub_masks(mask_image_open, w, h)
        for color, sub_mask in sub_masks.items():
            category_id = category_colors[color]

            # "annotations" info
            polygons, segmentations = create_sub_mask_annotation(sub_mask)

            # Check if we have classes that are a multipolygon
            if category_id in multipolygon_ids:
                # Combine the polygons to calculate the bounding box and area
                multi_poly = MultiPolygon(polygons)

                annotation = create_annotation_format(multi_poly, segmentations, image_id, category_id, annotation_id)

                annotations.append(annotation)
                annotation_id += 1
            else:
                for i in range(len(polygons)):
                    # Cleaner to recalculate this variable
                    segmentation = [np.array(polygons[i].exterior.coords).ravel().tolist()]

                    annotation = create_annotation_format(polygons[i], segmentation, image_id, category_id, annotation_id)

                    annotations.append(annotation)
                    annotation_id += 1
        image_id += 1
    return images, annotations, annotation_id

if __name__ == "__main__":
    # Get the standard COCO JSON format
    coco_format = get_coco_json_format()

    for keyword in ["train", "val"]:
        mask_path = "dataset/{}_mask/".format(keyword)

        # Create category section
        coco_format["categories"] = create_category_annotation(category_ids)

        # Create images and annotations sections
        coco_format["images"], coco_format["annotations"], annotation_cnt = images_annotations_info(mask_path)

        with open("output/{}.json".format(keyword),"w") as outfile:
            json.dump(coco_format, outfile)

        print("Created %d annotations for images in folder: %s" % (annotation_cnt, mask_path))

标签： datasetmaskcoco

dataset - 如何从黑白蒙版生成 COCO 数据集

问题描述

解决方案

推荐阅读