coco数据集生成-原本有mask标注的单类多bbox--转换为每一个bbox扩边两倍为单独一个图生成新的coco形式json文件

# -*- coding: utf-8 -*-
import sys, getopt
import os
import json
import cv2
import random
import numpy as np
np.random.seed(41)
import glob
import shutil
import os

# -*- coding: utf-8 -*-
import os
import sys, getopt
from pycocotools.coco import COCO, maskUtils
import cv2
import numpy as np


def mkdir_os(path):
    if not os.path.exists(path):
        os.makedirs(path)

def groberect(points, ww, hh):
    x1 = points[0]
    y1 = points[1]
    w = points[2]
    h = points[3]

    x2 = x1 + w - 1
    y2 = y1 + h - 1

    px = float(x1 + x2) / 2
    py = float(y1 + y2) / 2

    w = w * 2
    h = h * 2

    l = max(0, px - w / 2)
    r = min(ww - 1, px + w / 2) #x2
    t = max(0, py - h / 2)
    b = min(hh - 1, py + h / 2) #y2

    w = r - l + 1
    h = b - t + 1

    # x1y1 x2y2
    return [int(l), int(t), int(w), int(h)]

def main(argv):
    # json_file = './data/coco/annotations/instances_val2017.json'
    # dataset_dir = './data/coco/val2017/'
    # save_dir = './data/coco/vis/'

    inputfile = ''
    jsonfile = ''
    outputfile = ''

    try:
        opts, args = getopt.getopt(argv, "hi:j:o:", ["ifile=", "jfile=", "ofile="])
    except getopt.GetoptError:
        print('test.py -i <inputfile> -j <jsonfile> -o <outputfile>')
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-h':
            print('test.py -i <inputfile> -j <jsonfile> -o <outputfile>')
            sys.exit()
        elif opt in ("-i", "--ifile"):
            inputfile = arg
        elif opt in ("-j", "--jfile"):
            jsonfile = arg
        elif opt in ("-o", "--ofile"):
            outputfile = arg

    print('\n输入的文件为：', inputfile)
    print('\n输入的json为：', jsonfile)
    print('\n输出的文件为：', outputfile)


    images = []
    annotations = []

    mkdir_os(outputfile)

    coco = COCO(jsonfile)
    catIds = coco.getCatIds(catNms=['wires'])  # catIds=1 表示人这一类
    imgIds = coco.getImgIds(catIds=catIds)  # 图片id，许多值

    image_id = 0
    annotations_id = 0

    for i in range(len(imgIds)):
        if i % 100 == 0:
            print(i, "/", len(imgIds))
        img_info = coco.loadImgs(imgIds[i])[0]

        cvImage = cv2.imread(os.path.join(inputfile, img_info['file_name']), -1)
        cvImage = cv2.cvtColor(cvImage, cv2.COLOR_BGR2GRAY)
        cvImage = cv2.cvtColor(cvImage, cv2.COLOR_GRAY2BGR)

        ori_H, ori_W = cvImage.shape[:2]

        annIds = coco.getAnnIds(imgIds=img_info['id'], catIds=catIds, iscrowd=None)
        anns = coco.loadAnns(annIds)


        #原来一幅图可能拆成多个，需要保存多个图像，生成多个image信息，多个annotation信息
        img_info_append = []
        for index in range(len(anns)):

            ann = anns[index]
            img_temp = img_info.copy()
            if 'segmentation' in ann: #只处理存在annotation的情况
                if type(ann['segmentation']) == list: #只处理points这种形式的mask标注的情况
                    bbox = ann['bbox']
                    new_bbox = groberect(bbox, ori_W, ori_H)

                    img_temp['width'] = new_bbox[2]
                    img_temp['height'] = new_bbox[3]
                    img_temp['file_name'] = img_temp['file_name'].split('.')[0] + "_" + str(index) + ".jpg"
                    img_temp['id'] = image_id

                    #cropImg = img[(y):(y + hh), (x):(x + ww)]
                    save_cvImage = cvImage[new_bbox[1]:(new_bbox[1] + new_bbox[3]), new_bbox[0]:(new_bbox[0] + new_bbox[2])]
                    cv2.imwrite(os.path.join(outputfile, img_temp['file_name']), save_cvImage)

                    img_info_append.append(img_temp)

                    ann['bbox'] = [bbox[0]-new_bbox[0], bbox[1]-new_bbox[1], bbox[2], bbox[3]]
                    for seg in range(len(ann['segmentation'])):
                        ori_seg = ann['segmentation'][seg]
                        for k in range(len(ori_seg)//2):
                            ori_seg[2 * k] = ori_seg[2 * k] - new_bbox[0]
                            ori_seg[2 * k + 1] = ori_seg[2 * k + 1] - new_bbox[1]
                        ann['segmentation'][seg] = ori_seg

                    ann['image_id'] = int(image_id)
                    #anns[index] = ann
                else:
                    exit()
            else:
                exit()

            image_id += 1
            #annotations_id += 1

        annotations.extend(anns)
        images.extend(img_info_append)
    instance = {}
    instance['info'] = 'spytensor created'
    instance['license'] = ['license']
    instance['images'] = images
    instance['annotations'] = annotations
    instance['categories'] = coco.dataset['categories']

    import io
    #json.dump(instance, io.open("./result.json", 'w', encoding='utf-8'), ensure_ascii=False, indent=1)  # indent=2 更加美观显示
    with io.open("./new_json.json", 'w', encoding="utf-8") as outfile:
        my_json_str = json.dumps(instance, ensure_ascii=False, indent=1)
        #python3 无
        # if isinstance(my_json_str, str):
        #     my_json_str = my_json_str.decode("utf-8")
        outfile.write(my_json_str)


if __name__ == "__main__":
    main(sys.argv[1:])