如何从黑白蒙版生成 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,# Outlier
    "(255,0)": 1,# Window
    "(255,255,0)": 2,# Wall
    "(128,255)": 3,# Balcony
    "(255,128,0)": 4,# Door
    "(0,255)": 5,# Roof
    "(128,255)": 6,# Sky
    "(0,0)": 7,# Shop
    "(128,128)": 8,# Chimney
    "(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,h,image_id)
        images.append(image)

        sub_masks = create_sub_masks(mask_image_open,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,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"],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))

解决方法

暂无找到可以解决该程序问题的有效方法,小编努力寻找整理中!

如果你已经找到好的解决方法,欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@)