【亲测有用】使用Image和Mask图像转化成COCO数据格式

方法1: 轮廓处理

import json
import numpy as np
from pycocotools import mask
from skimage import measure
import cv2
import os
import sys
if sys.version_info[0] >= 3:
    unicode = str
__author__ = 'hcaesar'
import io
def maskToanno(ground_truth_binary_mask,ann_count,category_id):
    fortran_ground_truth_binary_mask = np.asfortranarray(ground_truth_binary_mask)
    encoded_ground_truth = mask.encode(fortran_ground_truth_binary_mask)
    ground_truth_area = mask.area(encoded_ground_truth)
    ground_truth_bounding_box = mask.toBbox(encoded_ground_truth)
    contours = measure.find_contours(ground_truth_binary_mask, 0.5)

    annotation = {
        "segmentation": [],
        "area": ground_truth_area.tolist(),
        "iscrowd": 0,
        "image_id": ann_count,
        "bbox": ground_truth_bounding_box.tolist(),
        "category_id": category_id,
        "id": ann_count
    }
    for contour in contours:
        contour = np.flip(contour, axis=1)
        segmentation = contour.ravel().tolist()
        annotation["segmentation"].append(segmentation)
    return annotation
block_mask_path="D:\\6Ddataset\\XieSegmentation\\1block\\block_mask_thresh"
mouse_mask_path="D:\\6Ddataset\\XieSegmentation\\2mouse\\mouse_mask_thresh"
block_mask_image_files=os.listdir(block_mask_path)
mouse_mask_image_files=os.listdir(mouse_mask_path)
jsonPath="blockmouseAll.json"
annCount=0
imageCount=0
path="D:\\6Ddataset\\XieSegmentation\\1block\\block_rgb"
rgb_image_files=os.listdir(path)
with io.open(jsonPath, 'w', encoding='utf8') as output:
    # 那就全部写在一个文件夹好了
    # 先写images的信息
    output.write(unicode('{\n'))
    output.write(unicode('"images": [\n'))
    for image in rgb_image_files:
        output.write(unicode('{'))
        annotation = {
            "height": 480,
            "width": 640,
            "id": imageCount,
            "file_name": image
        }
        str_ = json.dumps(annotation,indent=4)
        str_ = str_[1:-1]
        if len(str_) > 0:
            output.write(unicode(str_))
            imageCount = imageCount + 1
        if(image=="0000000599_rgb.png"):
            output.write(unicode('}\n'))
        else:
            output.write(unicode('},\n'))
    output.write(unicode('],\n'))
    #接下来写cate
    output.write(unicode('"categories": [\n'))
    output.write(unicode('{\n'))
    categories={
        "supercategory": "block",
        "id": 0,
        "name": "block"
    }
    str_ = json.dumps(categories, indent=4)
    str_ = str_[1:-1]
    if len(str_) > 0:
        output.write(unicode(str_))
    output.write(unicode('},\n'))
    output.write(unicode('{\n'))
    categories={
        "supercategory": "mouse",
        "id": 1,
        "name": "mouse"
    }
    str_ = json.dumps(categories, indent=4)
    str_ = str_[1:-1]
    if len(str_) > 0:
        output.write(unicode(str_))
    output.write(unicode('}\n'))
    output.write(unicode('],\n'))
    output.write(unicode('"annotations": [\n'))
    for i in range(len(block_mask_image_files)):
    #for (block_image,mouse_image) in (block_mask_image_files,mouse_mask_image_files):
        block_image=block_mask_image_files[i]
        mouse_image=mouse_mask_image_files[i]
        #output.write(unicode('{\n'))
        block_im=cv2.imread(os.path.join(block_mask_path,block_image),0)
        mouse_im = cv2.imread(os.path.join(mouse_mask_path, mouse_image), 0)
        _,block_im=cv2.threshold(block_im,100,1,cv2.THRESH_BINARY)
        _, mouse_im = cv2.threshold(mouse_im, 100, 1, cv2.THRESH_BINARY)
        block_im=np.array(block_im).astype(np.uint8)
        mouse_im = np.array(mouse_im).astype(np.uint8)
        block_anno=maskToanno(block_im,annCount,0)
        mouse_anno = maskToanno(mouse_im, annCount,1)
        str_block = json.dumps(block_anno,indent=4)
        str_block = str_block[1:-1]
        str_mouse = json.dumps(mouse_anno,indent=4)
        str_mouse = str_mouse[1:-1]
        if len(str_block) > 0:
            output.write(unicode('{\n'))
            output.write(unicode(str_block))
            output.write(unicode('},\n'))
            output.write(unicode('{\n'))
            output.write(unicode(str_mouse))
            if (block_image == "0000000599_labels.png"):
                output.write(unicode('}\n'))
            else:
                output.write(unicode('},\n'))
            #output.write(unicode('},\n'))
            annCount = annCount + 1
    output.write(unicode(']\n'))
    output.write(unicode('}\n'))
        #output.write(unicode('},\n'))
#cate

方法2-利用pycococreator

import cv2
import numpy as np
import os, glob
 
 
def rgb2masks(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
    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 = lbl * mask[..., None]      # colorful leaf with black background
            # np.repeat(mask[...,None],3,axis=2)    # 3D mask
            leaf = np.where(mask[..., None], white_mask, 0)
            mask_name = './shapes/train/annotations/' + lbl_id + '_leaf_' + str(idx) + '.png'
            cv2.imwrite(mask_name, leaf)
            idx += 1
 
 
label_dir = './labels'
label_list = glob.glob(os.path.join(label_dir, '*.png'))
for label_name in label_list:
    rgb2masks(label_name)

利用pycococreator和得到的黑白masks生成coco json格式的数据集，代码如下:

import datetime
import json
import os
import re
import fnmatch
from PIL import Image
import numpy as np
from pycococreatortools import pycococreatortools
 
 
ROOT_DIR = './shapes/train'
IMAGE_DIR = os.path.join(ROOT_DIR, "shapes_train2017")
ANNOTATION_DIR = os.path.join(ROOT_DIR, "annotations")
 
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():
    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 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('{}/instances_leaf_train2017.json'.format(ROOT_DIR), 'w') as output_json_file:
        json.dump(coco_output, output_json_file)
 
 
if __name__ == "__main__":
    main()