将mask图像转换成coco形式（包含bbox）

最近在医学图像上visual grounding的实验，由于大多息肉的任务都是分割，我之前收集的数据集也是无监督数据集or分割数据集，因此萌发了将mask转成bounding box标签的念头，来扩充目标检测数据集。

借鉴以下博客，但是我的数据在他的代码上有问题，所以做了一些修改。数据集：息肉二分类分割数据集，一张图中可能包含多个息肉。【coco】掩膜mask影像转coco格式txt（含python代码）_实例分割 将mask转化为txt-CSDN博客

coco数据集的格式

一张图片可以对应多个annotations，所以id是累加的，但注意images中的id要与annotations中的image_id一致。

{
"info": [
{

    "year": "2023",
    "version": "1",
    "contributor": "bulibuli",
    "url": "",
    "date_created": "2023-01-17"
}
],
"lisence": [
{

    "id": 1,
    "url": "https://creativecommons.org/licenses/by/4.0/",
    "name": "CC BY 4.0"
}
],
"categories": [
{

    "supercategory": "polyp",
    "id": 1,
    "name": "polyp"
}
],
"images": [
{
    "height": 500,
    "width": 574,
    "id": 1,
    "file_name": "CVC-300151.png"
},
...
"annotations": [
{

    "segmentation": [
        [
            446,
            517,
            447,
            517,
            448,
            517,
            449,
            517,
            450,
            517,
            451,
            517,
            452,
            517,
            453,
            517,
            454,
            517,
            455,
            517
        ]
    ],
    "area": 46450.0,
    "iscrowd": 0,
    "image_id": 1,
    "bbox": [
        776,
        663,
        324,
        246
    ],
    "category_id": 1,
    "id": 1
},
.....
]
}

路径设置

图像和mask的文件名必须一一对应，一模一样。

# mask图像路径
block_mask_path = 'your path'
block_mask_image_files = sorted(os.listdir(block_mask_path))

 
# coco json保存的位置
jsonPath = "your path/label.json"
annCount = 1
imageCount = 1
# 原图像的路径， 原图像和mask图像的名称是一致的。
path = "your path"
rgb_image_files = sorted(os.listdir(path))
if block_mask_image_files!=rgb_image_files: print("error")

写入基本信息和lisence

with io.open(jsonPath, 'w', encoding='utf8') as output:
    # 那就全部写在一个文件夹好了
    output.write(unicode('{\n'))
    # 基本信息
    output.write(unicode('"info": [\n'))
    output.write(unicode('{\n'))
    info={
        "year": "2023",
        "version": "1",
        "contributor": "bulibuli",
        "url": "",
        "date_created": "2023-01-17"
    }
    str_ = json.dumps(info, indent=4)
    str_ = str_[1:-1]
    if len(str_) > 0:
        output.write(unicode(str_))
    output.write(unicode('}\n'))
    output.write(unicode('],\n'))

    #lisence
    output.write(unicode('"lisence": [\n'))
    output.write(unicode('{\n'))
    info={
        "id": 1,
        "url": "https://creativecommons.org/licenses/by/4.0/",
        "name": "CC BY 4.0"
    }
    str_ = json.dumps(info, indent=4)
    str_ = str_[1:-1]
    if len(str_) > 0:
        output.write(unicode(str_))
    output.write(unicode('}\n'))
    output.write(unicode('],\n'))

写入类别

这里只有一个类别：息肉，因此直接写入即可，如果有多个类别，也可以往里面添加，请自助修改，但是一般多分类的数据集应该不会只有一个mask吧。

# category
    output.write(unicode('"categories": [\n'))
    output.write(unicode('{\n'))
    categories = {
        "supercategory": "polyp",
        "id": 1,
        "name": "polyp"
    }
    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'))

写入图像信息

需要用cv2读入图像，并获取宽高。

# images    
    output.write(unicode('"images": [\n'))
    for image in rgb_image_files:
        if os.path.exists(os.path.join(block_mask_path, image)):
            output.write(unicode('{'))
            block_im = cv2.imread(os.path.join(path, image))
            h,w,_=block_im.shape
            annotation = {
                "height": h,
                "width": w,
                "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 == rgb_image_files[-1]):
                output.write(unicode('}\n'))
            else:
                output.write(unicode('},\n'))
    output.write(unicode('],\n'))

写入标签信息

读取二值图像并转成array，获取该图像的annotations，传入image_id，类别id。对于每个annotation写入文件。

# 写annotations
    output.write(unicode('"annotations": [\n'))
    for i in range(len(block_mask_image_files)):
        if os.path.exists(os.path.join(path, block_mask_image_files[i])):
            block_image = block_mask_image_files[i]
            # print(block_image)
            # 读取二值图像
            block_im = cv2.imread(os.path.join(block_mask_path, block_image), 0)
            _, block_im = cv2.threshold(block_im, 100, 1, cv2.THRESH_BINARY)
            if not block_im is None:
                block_im = np.array(block_im, dtype=object).astype(np.uint8)
                block_anno = maskToanno(block_im, annCount, 1)
                # print(block_image,len(block_anno))
                for b in block_anno:
                    str_block = json.dumps(b, indent=4)
                    str_block = str_block[1:-1]
                    if len(str_block) > 0:
                        output.write(unicode('{\n'))
                        output.write(unicode(str_block))
                        if (block_image == rgb_image_files[-1] and b == block_anno[-1]):
                            output.write(unicode('}\n'))
                        else:
                            output.write(unicode('},\n'))
                annCount = annCount + 1
            else:
                print(block_image)

masktoanno函数

首先利用cv2.findContours函数，找到所有的轮廓，这里只取外轮廓，因为我们的分割mask是实心的。该函数有两个输出，我们取第一个输出，是一个列表存储了所有轮廓，如果该列表的长度为0，则说明， mask图像是全黑的，该图像中没有息肉。

然后对于每一个轮廓，如果长度小于3，则说明该轮廓构不成一个面，因此该轮廓作废。其余的轮廓我们直接采取cv2.contourArea函数获取mask的面积，然后通过cv2.boundingRect直接将轮廓转成bounding box，添加进annotation，加入文件。

global segmentation_id
segmentation_id = 1
# annotations部分的实现
def maskToanno(ground_truth_binary_mask, ann_count, category_id):
    contours, _ = cv2.findContours(ground_truth_binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)  # 根据二值图找轮廓
    annotations = [] #一幅图片所有的annotatons
    # print(len(contours),contours)
    global segmentation_id
    if(len(contours)==0):print("0")
    # 对每个实例进行处理
    for i,contour in enumerate(contours):
        if(len(contour)<3):
            print("The contour does not constitute an area")
            continue
        ground_truth_area = cv2.contourArea(contour)
        x, y, w, h = cv2.boundingRect(contour)
        annotation = {
            "segmentation": [],
            "area": ground_truth_area,
            "iscrowd": 0,
            "image_id": ann_count,
            "bbox": [x,y,w,h],
            "category_id": category_id,
            "id": segmentation_id
        }
        # 求segmentation部分
        contour = np.flip(contour, axis=0)
        segmentation = contour.ravel().tolist()
        annotation["segmentation"].append(segmentation)
        annotations.append(annotation)
        segmentation_id = segmentation_id + 1
    return annotations

总代码

目前生成的json没有问题，但是还没训练看效果，效果过几天再更新。

import json
import numpy as np
from pycocotools import mask
import cv2
import os
import sys
 
if sys.version_info[0] >= 3:
    unicode = str
 
 
import io
# 实例的id，每个图像有多个物体每个物体的唯一id
global segmentation_id
segmentation_id = 1
# annotations部分的实现
def maskToanno(ground_truth_binary_mask, ann_count, category_id):
    contours, _ = cv2.findContours(ground_truth_binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)  # 根据二值图找轮廓
    annotations = [] #一幅图片所有的annotatons
    # print(len(contours),contours)
    global segmentation_id
    if(len(contours)==0):print("0")
    # 对每个实例进行处理
    for i,contour in enumerate(contours):
        if(len(contour)<3):
            print("The contour does not constitute an area")
            continue
        ground_truth_area = cv2.contourArea(contour)
        x, y, w, h = cv2.boundingRect(contour)
        annotation = {
            "segmentation": [],
            "area": ground_truth_area,
            "iscrowd": 0,
            "image_id": ann_count,
            "bbox": [x,y,w,h],
            "category_id": category_id,
            "id": segmentation_id
        }
        # 求segmentation部分
        contour = np.flip(contour, axis=0)
        segmentation = contour.ravel().tolist()
        annotation["segmentation"].append(segmentation)
        annotations.append(annotation)
        segmentation_id = segmentation_id + 1
    return annotations
 
# mask图像路径
block_mask_path = '...'
block_mask_image_files = sorted(os.listdir(block_mask_path))

 
# coco json保存的位置
jsonPath = "..."
annCount = 1
imageCount = 1
# 原图像的路径， 原图像和mask图像的名称是一致的。
path = "..."
rgb_image_files = sorted(os.listdir(path))
if block_mask_image_files!=rgb_image_files: print("error")

with io.open(jsonPath, 'w', encoding='utf8') as output:
    # 那就全部写在一个文件夹好了
    output.write(unicode('{\n'))
    # 基本信息
    output.write(unicode('"info": [\n'))
    output.write(unicode('{\n'))
    info={
        "year": "2023",
        "version": "1",
        "contributor": "bulibuli",
        "url": "",
        "date_created": "2023-01-17"
    }
    str_ = json.dumps(info, indent=4)
    str_ = str_[1:-1]
    if len(str_) > 0:
        output.write(unicode(str_))
    output.write(unicode('}\n'))
    output.write(unicode('],\n'))

    #lisence
    output.write(unicode('"lisence": [\n'))
    output.write(unicode('{\n'))
    info={
        "id": 1,
        "url": "https://creativecommons.org/licenses/by/4.0/",
        "name": "CC BY 4.0"
    }
    str_ = json.dumps(info, indent=4)
    str_ = str_[1:-1]
    if len(str_) > 0:
        output.write(unicode(str_))
    output.write(unicode('}\n'))
    output.write(unicode('],\n'))

    # category
    output.write(unicode('"categories": [\n'))
    output.write(unicode('{\n'))
    categories = {
        "supercategory": "polyp",
        "id": 1,
        "name": "polyp"
    }
    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'))

    # images

    output.write(unicode('"images": [\n'))
    for image in rgb_image_files:
        if os.path.exists(os.path.join(block_mask_path, image)):
            output.write(unicode('{'))
            block_im = cv2.imread(os.path.join(path, image))
            h,w,_=block_im.shape
            annotation = {
                "height": h,
                "width": w,
                "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 == rgb_image_files[-1]):
                output.write(unicode('}\n'))
            else:
                output.write(unicode('},\n'))
    output.write(unicode('],\n'))
 
    # 写annotations
    output.write(unicode('"annotations": [\n'))
    for i in range(len(block_mask_image_files)):
        if os.path.exists(os.path.join(path, block_mask_image_files[i])):
            block_image = block_mask_image_files[i]
            # print(block_image)
            # 读取二值图像
            block_im = cv2.imread(os.path.join(block_mask_path, block_image), 0)
            _, block_im = cv2.threshold(block_im, 100, 1, cv2.THRESH_BINARY)
            if not block_im is None:
                block_im = np.array(block_im, dtype=object).astype(np.uint8)
                block_anno = maskToanno(block_im, annCount, 1)
                # print(block_image,len(block_anno))
                for b in block_anno:
                    str_block = json.dumps(b, indent=4)
                    str_block = str_block[1:-1]
                    if len(str_block) > 0:
                        output.write(unicode('{\n'))
                        output.write(unicode(str_block))
                        if (block_image == rgb_image_files[-1] and b == block_anno[-1]):
                            output.write(unicode('}\n'))
                        else:
                            output.write(unicode('},\n'))
                annCount = annCount + 1
            else:
                print(block_image)
            
    output.write(unicode(']\n'))
    output.write(unicode('}\n'))