coco数据集提取instance 和bbox

import os

import random,shutil

import cv2

from tqdm import tqdm

import numpy as np

from pycocotools.coco import COCO

from pycocotools import mask as coco_mask

from PIL import Image, ImageDraw

import matplotlib.pyplot as plt

import imgviz

# annotation_file = 'datasets/coco2017_zip/annotations/instances_val2017.json'

# coco = COCO(annotation_file)

# catIds = coco.getCatIds()

# imgIds = coco.getImgIds()

# print(catIds)

headstr = """\

<annotation>

<folder>VOC</folder>

<filename>%s</filename>

<source>

<database>My Database</database>

<annotation>COCO</annotation>

<image>flickr</image>

<flickrid>NULL</flickrid>

</source>

<owner>

<flickrid>NULL</flickrid>

<name>company</name>

</owner>

<size>

<width>%d</width>

<height>%d</height>

<depth>%d</depth>

</size>

<segmented>0</segmented>

"""

objstr = """\

<object>

<name>%s</name>

<pose>Unspecified</pose>

<truncated>0</truncated>

<difficult>0</difficult>

<bndbox>

<xmin>%d</xmin>

<ymin>%d</ymin>

<xmax>%d</xmax>

<ymax>%d</ymax>

</bndbox>

</object>

"""

tailstr = '''\

</annotation>

'''

def mkr(path):

if not os.path.exists(path):

os.makedirs(path) # 可以创建多级目录

def id2name(coco):

classes=dict()

for cls in coco.dataset['categories']:

classes[cls['id']]=cls['name']

return classes

def save_annotations_and_imgs(coco,dataset,filename,objs):

#将图片转为xml，例:COCO_train2017_000000196610.jpg-->COCO_train2017_000000196610.xml

dst_anno_dir = os.path.join(anno_dir, dataset)

mkr(dst_anno_dir)

anno_path=dst_anno_dir + '/' + filename[:-3]+'xml'

img_path=dataDir+dataset+'/'+filename

print("img_path: ", img_path)

dst_img_dir = os.path.join(img_dir, dataset)

mkr(dst_img_dir)

dst_imgpath=dst_img_dir+ '/' + filename

print("dst_imgpath: ", dst_imgpath)

img=cv2.imread(img_path)

#if (img.shape[2] == 1):

# print(filename + " not a RGB image")

# return

shutil.copy(img_path, dst_imgpath)

head=headstr % (filename, img.shape[1], img.shape[0], img.shape[2])

tail = tailstr

# write_xml(anno_path,head, objs, tail)

def showimg(coco,dataset,img,classes,cls_id,show=True):

global dataDir

I=Image.open('%s/%s/%s'%(dataDir,dataset,img['file_name']))

#通过id，得到注释的信息

annIds = coco.getAnnIds(imgIds=img['id'], catIds=cls_id, iscrowd=None)

# print(annIds)

anns = coco.loadAnns(annIds)

# print(anns)

# coco.showAnns(anns)

objs = []

for ann in anns:

class_name=classes[ann['category_id']]

if class_name in classes_names:

print(class_name)

if 'bbox' in ann:

bbox=ann['bbox']

xmin = int(bbox[0])

ymin = int(bbox[1])

xmax = int(bbox[2] + bbox[0])

ymax = int(bbox[3] + bbox[1])

obj = [class_name, xmin, ymin, xmax, ymax]

objs.append(obj)

draw = ImageDraw.Draw(I)

draw.rectangle([xmin, ymin, xmax, ymax])

def save_colored_mask(mask, save_path):

lbl_pil = Image.fromarray(mask.astype(np.uint8), mode="P")

colormap = imgviz.label_colormap()

lbl_pil.putpalette(colormap.flatten())

lbl_pil.save(save_path)

if __name__ == '__main__':

savepath="datasets/coco2017_zip/my_coco/train2017/"

img_dir=savepath+'images'

anno_dir=savepath+'annotations'

mask_dir = savepath + 'mask_small'

datasets_list=['train2017']#['train2017', 'val2017'] #

#coco有80类，这里写要提取类的名字，以person为例

classes_dict = {'bottle':44,'chair':62,'couch':63,'potted plant':64,'dining table':67,'book':84,'bed':65}

classes_names = ['bottle','chair','couch','potted plant','dining table','book','bed']

# classes_names = ['bottle']

dataDir= 'datasets/coco2017_zip/'

for dataset in datasets_list:

#./COCO/annotations/instances_train2017.json

annFile='{}/annotations/instances_{}.json'.format(dataDir,dataset)

#使用COCO API用来初始化注释数据

coco = COCO(annFile)

#获取COCO数据集中的所有类别

classes = id2name(coco)

print(classes)

#[1, 2, 3, 4, 6, 8]

classes_ids = coco.getCatIds(catNms=classes_names)

print(classes_ids)

for cls in classes_names:

#获取该类的id

cls_id=coco.getCatIds(catNms=[cls])

img_ids=coco.getImgIds(catIds=cls_id)

print(cls,len(img_ids))

# imgIds=img_ids[0:10]

for imgId in tqdm(img_ids):

img = coco.loadImgs(imgId)[0]

filename = img['file_name']

img_path=dataDir+dataset+'/'+filename

image = cv2.imread(img_path)

# print(filename)

# objs=showimg(coco, dataset, img, classes,classes_ids,show=False)

annIds = coco.getAnnIds(imgIds=img['id'], catIds=classes_ids, iscrowd=None)

if len(annIds)==0:

print(annIds)

# print('annIds:',len(annIds))

anns = coco.loadAnns(annIds)

bboxes = []

for ann in anns:

class_name=classes[ann['category_id']]

if class_name in classes_names:

# print(class_name)

if 'bbox' in ann:

bbox=ann['bbox']

xmin = int(bbox[0])

ymin = int(bbox[1])

xmax = int(bbox[2] + bbox[0])

ymax = int(bbox[3] + bbox[1])

bbox_ = [class_name, xmin, ymin, xmax, ymax]

bboxes.append(bbox_)

save_path = os.path.join(mask_dir,class_name)

if not os.path.exists(save_path):

os.mkdir(save_path)

# draw 2d bboxes

# cv2.rectangle(image, (xmin,ymin), (xmax,ymax), (255, 0, 255))

# get instance masks, and save each class instance into diffirent file path

# print('333',anns[0])

# exit()

mask = coco.annToMask(anns[0]) * anns[0]['category_id']

img_ = image.copy()

class_num = classes_dict[class_name]

img_[mask!=class_num] = 255

crop = img_[ymin:ymax,xmin:xmax]

# print([xmin,xmax,ymin,ymax])

new_filename = filename.split('.')[0]+'.png'

if os.path.isfile(os.path.join(save_path,new_filename)):

continue

else:

cv2.imwrite('{}/{}'.format(save_path,new_filename),crop)

# cv2.imwrite('datasets/coco2017_zip/my_coco/001.jpg',image)

# # get the whole one picture's instance mask

# mask = coco.annToMask(anns[0]) * anns[0]['category_id']

# for i in range(len(anns) - 1):

# mask += coco.annToMask(anns[i + 1]) * anns[i + 1]['category_id']

# save_colored_mask(mask,'{}/{}'.format('/home/yaoxing/workspace/datasets/coco2017_zip/my_coco/mask_whole',filename.split('.')[0]+'.png'))

# cv2.imwrite('{}/{}'.format('datasets/coco2017_zip/my_coco/images',filename),image)

# print(filename)

# exit()