我有一个由焊缝和掩码组成的数据集(焊缝为白色,背景为黑色),尽管我需要使用 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))