COCO數據集轉mask

書接上文，先馬克一下，之后再改

# -*- coding: utf-8 -*-
"""
Created on Wed Jul  1 14:45:07 2020

@author: mhshao
"""
from pycocotools.coco import COCO
import os
import shutil
from tqdm import tqdm
import matplotlib.pyplot as plt
import cv2
from PIL import Image, ImageDraw
import skimage.io as io
import json
import numpy as np
'''
路徑參數
'''
#原coco數據集的路徑
dataDir= 'newdata/'
#用於保存新生成的mask數據的路徑
savepath = "newdata/"

'''
數據集參數
'''
#coco有80類，這里寫要進行二值化的類的名字
#其他沒寫的會被當做背景變成黑色
#如我只需要car、bus、truck這三類數據
classes_names = ['car','bus','truck']  
#要處理的數據集，比如val2017、train2017等
#不建議多個數據集在一個list中
#一次提取一個數據集安全點_(:3」∠❀)_
datasets_list = ['val2017']

#生成保存路徑，函數抄的(›´ω`‹ )
#if the dir is not exists,make it,else delete it
def mkr(path):
    if os.path.exists(path):
        shutil.rmtree(path)
        os.mkdir(path)
    else:
        os.mkdir(path)
        
#生成mask圖
def mask_generator(coco,width,height,anns_list):
    mask_pic = np.zeros((height, width))
    #生成mask
    for single in anns_list:
            mask_single = coco.annToMask(single)
            mask_pic += mask_single
    #轉化為255
    for row in range(height):
            for col in range(width):
                if (mask_pic[row][col] > 0):
                    mask_pic[row][col] = 255
    mask_pic = mask_pic.astype(int)
    '''
    #轉為三通道
    imgs = np.zeros(shape=(height, width, 3), dtype=np.float32)
    imgs[:, :, 0] = mask_pic[:, :]
    imgs[:, :, 1] = mask_pic[:, :]
    imgs[:, :, 2] = mask_pic[:, :]
    imgs = imgs.astype(int)
    '''    
    return mask_pic
    
#處理json數據並保存二值mask
def get_mask_data(annFile,mask_to_save):
    #獲取COCO_json的數據
    coco = COCO(annFile)
    #拿到所有需要的圖片數據的id
    classes_ids = coco.getCatIds(catNms = classes_names)
    #取所有類別的並集的所有圖片id
    #如果想要交集，不需要循環，直接把所有類別作為參數輸入，即可得到所有類別都包含的圖片
    imgIds_list = []
    for idx in classes_ids:
        imgidx = coco.getImgIds(catIds=idx)
        imgIds_list += imgidx
    #去除重復的圖片
    imgIds_list = list(set(imgIds_list))
    
    #一次性獲取所有圖像的信息
    image_info_list = coco.loadImgs(imgIds_list)
    
    #對每張圖片生成一個mask
    for imageinfo in image_info_list:
        #獲取對應類別的分割信息
        annIds = coco.getAnnIds(imgIds = imageinfo['id'], catIds = classes_ids, iscrowd=None)
        anns_list = coco.loadAnns(annIds)
        #生成二值mask圖
        mask_image = mask_generator(coco,imageinfo['width'],imageinfo['height'],anns_list)
        #保存圖片
        file_name = mask_to_save + '/' +imageinfo['file_name'][:-4]+'.png'
        plt.imsave(file_name , mask_imageif __name__ == '__main__':
    #按單個數據集進行處理
    for dataset in datasets_list:
        #用來保存最后生成的mask圖像目錄
        mask_to_save = savepath + 'masks/' + dataset
        mkr(savepath + 'masks/')
        #生成路徑
        mkr(mask_to_save)

        #獲取要處理的json文件路徑
        #我這里用了之前自己生成的部分類別json
        #具體方法見我前一篇博客
        annFile='{}/annotations/instances_{}_sub.json'.format(dataDir,dataset)
        #處理數據
        get_mask_data(annFile,mask_to_save)
        print('Got all the masks of {} from {} ٩( ๑╹ ꇴ ╹)۶'.format(classes_names,dataset))

 

 

 

 

 

000000001532.png

 

 

 

 

 

000000097924.png

 

 

 

 

000000121242.png