本博客主要参考了:将彩色RGB分割标注图像数据集转换为COCO格式的json文件_Liu, Xu的博客-CSDN博客_json转rgb
1、将彩色图片转换为黑白mask
import cv2
import numpy as np
import os, glob
"""
python 将彩色图片中的每一个实例切分出来
"""
def makedir(new_dir):
if not os.path.exists(new_dir):
os.makedirs(new_dir)
def rgb2masks(annotations_dir, label_name):
lbl_id = os.path.split(label_name)[-1].split('.')[0]
lbl = cv2.imread(label_name, 1)
h, w = lbl.shape[:2]
leaf_dict = {}
idx = 0
white_mask = np.ones((h, w, 3), dtype=np.uint8) * 255
makedir(annotations_dir)
for i in range(h):
for j in range(w):
if tuple(lbl[i][j]) in leaf_dict or tuple(lbl[i][j]) == (0, 0, 0):
continue
leaf_dict[tuple(lbl[i][j])] = idx
mask = (lbl == lbl[i][j]).all(-1)
leaf = np.where(mask[..., None], white_mask, 0)
mask_name = annotations_dir + "/" + lbl_id + '_leaf_' + str(
idx) + '.png'
cv2.imwrite(mask_name, leaf)
idx += 1
if __name__ == "__main__":
label_dir = 'A4/A4_train_label' # 需要处理的文件名
annotations_dir = "A4/A4_train_annotations"
label_list = glob.glob(os.path.join(label_dir, '*.png'))
for label_name in label_list:
rgb2masks(annotations_dir, label_name)
print(label_name + "has done")
print("finish")
2、生成COCO格式的Json文件
import datetime
import json
import os
import re
import fnmatch
from PIL import Image
import numpy as np
from pycococreatortools import pycococreatortools
INFO = {
"description": "Leaf Dataset",
"url": "https://github.com/waspinator/pycococreator",
"version": "0.1.0",
"year": 2017,
"contributor": "Francis_Liu",
"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': 'leaf',
'supercategory': 'leaf',
}
]
def filter_for_jpeg(root, files):
file_types = ['*.jpeg', '*.jpg', '*.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(ROOT_DIR, IMAGE_DIR, ANNOTATION_DIR, json_name):
coco_output = {
"info": INFO,
"licenses": LICENSES,
"categories": CATEGORIES,
"images": [],
"annotations": []
}
ROOT_DIR = ROOT_DIR
IMAGE_DIR = IMAGE_DIR
ANNOTATION_DIR = ANNOTATION_DIR
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 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 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(ROOT_DIR + "/" + json_name, 'w') as output_json_file:
json.dump(coco_output, output_json_file)
if __name__ == "__main__":
ROOT_DIR = 'A4'
IMAGE_DIR = os.path.join(ROOT_DIR, "A4_test")
ANNOTATION_DIR = os.path.join(ROOT_DIR, "A4_test_annotations")
json_name = "instances_test2017.json"
main(ROOT_DIR, IMAGE_DIR, ANNOTATION_DIR, json_name)