coco data set conversion mask file

from pycocotools.coco import COCO
import cv2
import numpy as np
import os
from pycocotools import coco


def get_contours(img):
    img = cv2.imread(img)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    kernel = np.ones((1, 5), np.uint8)
    ret, img1 = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
    binary = cv2.morphologyEx(img1, cv2.MORPH_CLOSE, kernel, anchor=(2, 0), iterations=5)
    contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # cv2.drawContours(img1, contours, -1, (0, 0, 134), 3)
    return contours, hierarchy


def ReadFromCOCO(path):
    # 初始化COCO数据加载器
    coco = COCO(path)
    # 获取所有图片ID
    img_ids = coco.getImgIds()
    # 创建空列表用于存储图片路径和标注信息
    img_paths_list = []
    labels_list = []
    # 遍历每个图片ID
    for img_id in img_ids:
        # 获取图片信息
        img_info = coco.loadImgs(img_id)[0]
        # 获取图片路径
        img_path = img_info['file_name']
        # 将图片路径添加到img_paths_list中
        img_paths = os.path.join(img_dir, img_path)
        img_paths_list.append(img_paths)

        # 获取图片对应的标注信息
        ann_ids = coco.getAnnIds(imgIds=img_id)
        annotations = coco.loadAnns(ann_ids)
        # 创建一个空列表用于存储当前图片的标注信息
        img_labels = []
        for ann in annotations:
            # 获取类别ID
            category_id = ann['category_id']
            # 获取多边形边界信息
            segmentation = ann['segmentation']
            # 将类别ID和多边形边界信息分别存储在一个小列表中，并添加到img_labels中
            img_labels.append([category_id, segmentation])
        # 将当前图片的标注信息添加到labels_list中
        labels_list.append(img_labels)
    return img_paths_list, labels_list


def getMaskFromList(image_path, label):
    # 获取指定索引的图片路径
    Img_ori, Img_Masks = [], []
    print(image_path)
    # 读取图片
    Img = cv2.imread(image_path)
    # # 获取指定索引的标注信息
    # print("labels:",label)
    # # 创建空列表，用于存储单目标掩膜图像
    count = 0
    # file1=os.path.join(out_dir,image_path.split("/")[-1].split(".")[0])
    # if not os.path.exists(file1): os.makedirs(file1)
    for lab in label:
        print("label:", label)
        category_id, segmentation = lab
        print(category_id, segmentation)
        if isinstance(segmentation, dict):
            img1 = Img.copy()

            # 多边形形状，使用RLE解码
            rle, m = segmentation["counts"], segmentation["size"]
            print(m)

            assert sum(rle) == m[0] * m[1]
            M = np.zeros(m[0] * m[1])
            N = len(rle)
            n = 0
            val = 1
            for pos in range(N):
                val = not val
                for c in range(rle[pos]):
                    M[n] = val
                    if val:
                        M[n] = 255
                    n += 1

            GEMFIELD = M.reshape(([m[0], m[1]]), order="F")
            # file1 = os.path.join(out_dir, image_path.split("/")[-1].split(".")[0])
            # if not os.path.exists(file1): os.makedirs(file1)
            # file2 = os.path.join(os.path.join(out_dir, image_path.split("/")[-1].split(".")[0]), str(category_id))
            # if not os.path.exists(file2): os.makedirs(file2)
            # cv2.imwrite(os.path.join(file2, image_path.split("/")[-1].split(".")[0] + "_" + str(count) + ".jpg"),
            #             GEMFIELD)
            Img_Masks.append([category_id, GEMFIELD])
            # cv2.imshow('111', GEMFIELD)
            # mypath=os.path.join(file2, image_path.split("/")[-1].split(".")[0] + "_" + str(count) + ".jpg")
            # contours,hierarchy=get_contours(mypath)
            # cv2.drawContours(img1, contours, -1, (255,0,0), 3)
            # cv2.imshow("111222",img1)
            # cv2.waitKey(0)

        # elif isinstance(segmentation, list):
        else:
            # 列表形式，表示多边形的轮廓点
            # print(len(segmentation))
            for seg in segmentation:
                seg = np.array(seg, dtype=np.int32).reshape(-1, 2)
                mask = np.zeros(Img.shape[:2], dtype=np.uint8)
                img1 = Img.copy()
                file2 = os.path.join(os.path.join(out_dir, image_path.split("/")[-1].split(".")[0]), str(category_id))
                if not os.path.exists(file2): os.makedirs(file2)
                # cv2.fillPoly(img1, [seg], color=255)
                count += 1
                # 填充多边形区域
                cv2.fillPoly(mask, [seg], color=255)
                # cv2.circle(img1, [seg], 5, color=123, thickness=2)
                # cv2.imwrite(os.path.join(file2,image_path.split("/")[-1].split(".")[0]+"_"+str(count)+".png"),mask)
                # cv2.imwrite(os.path.join(file2, image_path.split("/")[-1].split(".")[0] + "_" + str(count) + ".png"),
                #             img1)
                # cv2.imshow("12345",img1)
                # cv2.waitKey(0)
                Img_Masks.append([category_id, mask])

    return Img, Img_Masks,


# 一张图对应生成一张mask图
# def getMaskFromList1(image_path, labels):
#     # 获取指定索引的图片路径
#     # 读取图片
#     Img = cv2.imread(image_path)
#     # 创建空列表，用于存储单目标掩膜图像
#     Img_Masks = []
#     # mask = np.zeros(Img.shape[:2], dtype=np.uint8)
#     # 遍历标注信息，获取单目标掩膜
#     for label in labels:
#         category_id, segmentation = label
#         mask = np.zeros(Img.shape[:2], dtype=np.uint8)
#         # 判断segmentation的数据类型，如果是字典，则是多边形形状，否则为列表
#         if isinstance(segmentation, dict):
#             # 多边形形状，使用RLE解码
#             rle, m = segmentation["counts"], segmentation["size"]
#             M = np.zeros(m[0] * m[1], dtype=np.uint8)
#             N = len(rle)
#             n = 0
#             val = 255
#             for pos in range(N):
#                 val = not val
#                 for c in range(rle[pos]):
#                     M[n] = val
#                     n += 1
#             for idx in range(len(M)):
#                 if M[idx]==1:
#                     M[idx]=255
#             GEMFIELD = M.reshape(([m[0], m[1]]), order="F")
#             cv2.imwrite(os.path.join(out_dir, img_path.split("/")[-1].split(".")[0]+"_1.png"), GEMFIELD)
#
#         elif isinstance(segmentation, list):
#             # 列表形式，表示多边形的轮廓点
#             mask1=[]
#             for seg in segmentation:
#                 seg = np.array(seg, dtype=np.int32).reshape(-1, 2)
#                 # mask = np.zeros(Img.shape[:2], dtype=np.uint8)
#                 # 填充多边形区域
#                 cv2.fillPoly(mask, [seg], color=255)
#                 Img_Masks.append(mask)
#             cv2.imwrite(os.path.join(out_dir, img_path.split("/")[-1]), mask)
#
#     return Img, Img_Masks

def get_image_size(img,size):#size=[width,height]
    img = np.array(img)
    height, width = img.shape
    ratio = width / height
    mask_image = []
    if width > height:
        w = size[0]
        h = int(w / ratio)
        img1 = cv2.resize(img, (w, h))
        h1, w1 = img1.shape
        if h1 % 2 == 0:
            c1 = int((size[1] - h1) / 2)
            c2 = int((size[1]- h1) / 2)
            mask1 = np.zeros([c1, w], dtype=np.uint8)
            mask2 = np.zeros([c2, w], dtype=np.uint8)
            # img2=cv2.vconcat([mask1,img1,mask2])
            img2 = np.concatenate([mask1, img1, mask2], axis=0)
            mask_image.append(img2)
            print("*****", img2.shape)

        elif h1 % 2 == 1:
            c1 = int(((size[1] - h1) / 2) + 1)
            c2 = int((size[1] - h1) / 2)
            mask1 = np.zeros([c1, w], dtype=np.uint8)
            mask2 = np.zeros([c2, w], dtype=np.uint8)
            # img2 = cv2.vconcat([mask1, img1, mask2])
            img2 = np.concatenate([mask1, img1, mask2], axis=0)
            mask_image.append(img2)

    elif width < height:
        h = size[1]
        w = int(h * ratio)
        img1 = cv2.resize(img, (w, h))
        h1, w1 = img1.shape
        if w1 % 2 == 0:
            c1 = int((size[0] - w1) / 2)
            c2 = int((size[0] - w1) / 2)
            mask1 = np.zeros([h, c1], dtype=np.uint8)
            mask2 = np.zeros([h, c2], dtype=np.uint8)
            # img2 = cv2.hconcat([mask1, img1, mask2])
            img2 = np.concatenate([mask1, img1, mask2], axis=1)
            mask_image.append(img2)
            print(">>>>>>>>>", img2.shape)
            # cv2.imwrite("", img2)

        elif w1 % 2 == 1:
            c1 = int((size[0] - w1) / 2) + 1
            c2 = int((size[0] - w1) / 2)
            print("c1c2:", c1, c2)
            mask1 = np.zeros([h, c1], dtype=np.uint8)
            mask2 = np.zeros([h, c2], dtype=np.uint8)
            # img2 = cv2.hconcat([mask1, img1, mask2])
            img2 = np.concatenate([mask1, img1, mask2], axis=1)
            print(">>>>>>>>>", img2.shape)
            mask_image.append(img2)

    else:
        img1 = cv2.resize(img, (size[0], size[1]))
        mask_image.append(img1)

    return mask_image[0]

if __name__ == '__main__':
    data_path = '/media/tianhailong/新加卷/coco_seg/instances_train2017.json'
    img_dir = "/media/tianhailong/新加卷/coco_seg/train2017"
    out_dir = "/media/tianhailong/新加卷/coco_seg2"
    if not os.path.exists(out_dir): os.makedirs(out_dir)
    img_paths_list, labels_list = ReadFromCOCO(data_path)

    for idx, img_path in enumerate(img_paths_list):
        # print(img_path.split("/")[-1])
        Img, Img_Masks = getMaskFromList(img_path, labels_list[idx])
        file1 = os.path.join(out_dir, img_path.split("/")[-1].split(".")[0])
        if not os.path.exists(file1): os.makedirs(file1)
        for class_id,mask_img in Img_Masks:
            count=0
            print(mask_img.shape)
            mask_image=get_image_size(mask_img,(640,640))
            file2 = os.path.join(os.path.join(out_dir, img_path.split("/")[-1].split(".")[0]), str(class_id))
            if not os.path.exists(file2): os.makedirs(file2)
            count+=1
            cv2.imwrite(os.path.join(file2,img_path.split("/")[-1].split(".")[0]+ "_%s"%(count)+".png"),mask_image)