rle coco voc yolo label标签互相转换

rle2coco

import datetime
import json
import os
import re
import fnmatch
from PIL import Image
import numpy as np
from pycococreatortools import pycococreatortools
import pandas as pd

from skimage.data import imread
import matplotlib.pyplot as plt

dataset_train = '../positive_samples/'
csv_train = 	'../positive_example.csv'
IMAGE_DIR = dataset_train

df = pd.read_csv(csv_train )  										 # read csv file

INFO = {
    "description": "Kaggle Dataset",
    "url": "https://github.com/pascal1129",
    "version": "0.1.0",
    "year": 2018,
    "contributor": "pascal1129",
    "date_created": datetime.datetime.utcnow().isoformat(' ')
}

LICENSES = [
    {
        "id": 1,
        "name": "Attribution-NonCommercial-ShareAlike License",
        "url": "http://creativecommons.org/licenses/by-nc-sa/2.0/"
    }
]

CATEGORIES = [
    {
        'id': 1,
        'name': 'ship',
        'supercategory': 'ship',
    },
]

# ref: https://www.kaggle.com/paulorzp/run-length-encode-and-decode
# mask_rle(string) --> rle_decode() -->  np.ndarry(np.unit8)    
# shape: (height,width) , 1 - mask, 0 - background
def rle_decode(mask_rle, shape=(768, 768)):
    s = mask_rle.split()
    starts =  np.asarray(s[0::2], dtype=int)
    lengths = np.asarray(s[1::2], dtype=int)

    starts -= 1
    ends = starts + lengths
    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)
    for lo, hi in zip(starts, ends):
        img[lo:hi] = 1
    return img.reshape(shape).T  # Needed to align to RLE direction


def filter_for_jpeg(root, files):
    file_types = ['*.jpeg', '*.jpg']
    file_types = r'|'.join([fnmatch.translate(x) for x in file_types])
    files = [os.path.join(root, f) for f in files]
    files = [f for f in files if re.match(file_types, f)]
    
    return files

def save_bad_ann(image_name, mask, segmentation_id):
    img = imread(os.path.join(IMAGE_DIR, image_name))
    fig, axarr = plt.subplots(1, 3)
    axarr[0].axis('off')
    axarr[1].axis('off')
    axarr[2].axis('off')
    axarr[0].imshow(img)
    axarr[1].imshow(mask)
    axarr[2].imshow(img)
    axarr[2].imshow(mask, alpha=0.4)
    plt.tight_layout(h_pad=0.1, w_pad=0.1)
    if not os.path.exists('tmp'):
    	os.makedirs('tmp')
    plt.savefig( os.path.join('./tmp', image_name.split('.')[0] +'_' +str(segmentation_id) +'.png') )
    plt.close()

def main():
    # 最终放进json文件里的字典
    coco_output = {
        "info": INFO,
        "licenses": LICENSES,
        "categories": CATEGORIES,
        "images": [],   # 放一个空列表占位置，后面再append
        "annotations": []
    }

    image_id = 1
    segmentation_id = 1
    
    # 最外层的循环是图片，因为图片的基本信息需要共享
    # IMAGE_DIR路径下找到所有的图片
    for root, _, files in os.walk(IMAGE_DIR):
        image_paths = filter_for_jpeg(root, files)  # 图片文件地址
        num_of_image_files = len(image_paths)       # 图片个数

        # 遍历每一张图片
        for image_path in image_paths:
            # 提取图片信息
            image = Image.open(image_path)
            image_name = os.path.basename(image_path)   # 不需要具体的路径，只要图片文件名
            image_info = pycococreatortools.create_image_info(
                image_id, image_name, image.size)
            coco_output["images"].append(image_info)

            # 内层循环是mask，把每一张图片的mask搜索出来
            rle_masks = df.loc[df['ImageId'] == image_name, 'EncodedPixels'].tolist()
            num_of_rle_masks = len(rle_masks)

            for index in range(num_of_rle_masks):
                #print(rle_masks)
                binary_mask = rle_decode(rle_masks[index])
                class_id = 1    # 所有图片的类别都是1，ship
                category_info = {'id': class_id, 'is_crowd': 0}
                annotation_info = pycococreatortools.create_annotation_info(
                    segmentation_id, image_id, category_info, binary_mask,
                    image.size, tolerance=2)

                # 不是所有的标注都会被转换,低质量标注会被过滤掉
                # 正常的标注加入数据集，不好的标注保存供观察
                if annotation_info is not None:
                    coco_output["annotations"].append(annotation_info)
                else:
                    save_bad_ann(image_name, binary_mask, segmentation_id)
                    
                # 无论标注是否被写入数据集，均分配一个编号
                segmentation_id = segmentation_id + 1   

            print("%d of %d is done."%(image_id,num_of_image_files))
            image_id = image_id + 1

    with open('../instances_ships_train2018.json', 'w') as output_json_file:
        # json.dump(coco_output, output_json_file)
        json.dump(coco_output, output_json_file,indent=4)

if __name__ == "__main__":
    main()

coco2voc

from pycocotools.coco import COCO
import os
import shutil
from tqdm import tqdm
import skimage.io as io
import matplotlib.pyplot as plt
import cv2
from PIL import Image, ImageDraw
 
#the path you want to save your results for coco to voc
savepath="F:\\project\\DataSet\\"
img_dir=savepath+'samples/'
anno_dir=savepath+'Annotations/'

datasets_list=['train2018']
 
classes_names = ['ship']
#Store annotations and train2014/val2014/... in this folder
dataDir= img_dir
 
headstr = """\
<annotation>
    <folder>VOC</folder>
    <filename>%s</filename>
    <source>
        <database>Ship_Database</database>
        <annotation>COCO</annotation>
        <image>SHIP</image>
        <flickrid>NULL</flickrid>
    </source>
    <owner>
        <flickrid>NULL</flickrid>
        <name>d</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>
'''
 
#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)
#mkr(img_dir)
#mkr(anno_dir)
def id2name(coco):
    classes=dict()
    for cls in coco.dataset['categories']:
        classes[cls['id']]=cls['name']
    return classes
 
def write_xml(anno_path,head, objs, tail):
    f = open(anno_path, "w")
    f.write(head)
    for obj in objs:
        f.write(objstr%(obj[0],obj[1],obj[2],obj[3],obj[4]))
    f.write(tail)
 
 
def save_annotations_and_imgs(coco,dataset,filename,objs):
    #eg:COCO_train2014_000000196610.jpg-->COCO_train2014_000000196610.xml
    anno_path=anno_dir+filename[:-3]+'xml'
    #img_path=dataDir+dataset+'/'+filename
    img_path=dataDir+'/'+filename
    #print(img_path)
    dst_imgpath=img_dir+filename
 
    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'%(dataDir
                          ,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])
    if show:
        plt.figure()
        plt.axis('off')
        plt.imshow(I)
        plt.show()
 
    return objs
 
for dataset in datasets_list:
    #./COCO/annotations/instances_train2014.json
    annFile='{}/annotations/instances_{}.json'.format(dataDir,dataset)
    annFile=savepath+'instances_ships_train2018.json'
    #COCO API for initializing annotated data
    coco = COCO(annFile)
    '''
    COCO 对象创建完毕后会输出如下信息:
    loading annotations into memory...
    Done (t=0.81s)
    creating index...
    index created!
    至此, json 脚本解析完毕, 并且将图片和对应的标注数据关联起来.
    '''
    #show all classes in 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:
        #Get ID number of this class
        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']
            # print(filename)
            objs=showimg(coco, dataset, img, classes,classes_ids,show=False)
            #print(objs)
            save_annotations_and_imgs(coco, dataset, filename, objs)

voc2yolo

import xml.etree.ElementTree as ET
import os,sys
from os import listdir, getcwd
from os.path import join
import pickle
savepath="F:\\project\\DataSet\\"

anno_dir=savepath+'Annotations/labels\\'

classes=["0"]
def convert(size,box):
	# Conversion from VOC to YOLO format
	dw = 1./(size[0])
	dh = 1./(size[1])
	x = (box[0] + box[1])/2.0 - 1
	y = (box[2] + box[3])/2.0 - 1
	w = box[1] - box[0]
	h = box[3] - box[2]
	x = x*dw
	w = w*dw
	y = y*dh
	h = h*dh
	return (x,y,w,h)

#print ann_file
def converting_annotation(ann_file):
	for ann in ann_file:
		txt_file= ann.split('.')[0]+'.txt'
		in_file=open(anno_dir+ann)
		out_file =open(anno_dir+txt_file, 'w')
		tree=ET.parse(in_file)
		root = tree.getroot()
		size = root.find('size')
		w = int(size.find('width').text)
		h = int(size.find('height').text)

		print(w,h)
		for obj in root.iter('object'):
			cls=obj.find('name').text
			if cls not in classes:
				continue
			cls_id=classes.index(cls)
			xmlbox = obj.find('bndbox')
			b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
			data = convert((w,h), b)
			out_file.write(str(cls_id) + " " + " ".join([str(a) for a in data]) + '\n')


if __name__ == '__main__': 
	#ann_dir=sys.argv[1]

	ann_file=os.listdir(anno_dir)
	ann_file.sort()
	#print(ann_file)
	converting_annotation(ann_file)