VOC数据集生成代码使用说明

#split.py 文件 输入格式为images ,和标签txt文件，txt中的数据为坐标值共8个。

import os

import numpy as np

import math

import cv2 as cv

import imageio

#path = '/media/D/code/OCR/text-detection-ctpn/data/mlt_english+chinese/image'

#path = '/home/chendali1/Gsj/text-detection-ctpn-master/prepare_training_data/image/image_1000/'

path='/home/chendali1/Gsj/prepare_training_data/ICDAR/images_train/'

#gt_path = '/home/chendali1/Gsj/text-detection-ctpn-master/prepare_training_data/label/labelDigit1000/'

gt_path='/home/chendali1/Gsj/prepare_training_data/ICDAR/result_train/'

out_path = 're_image'

if not os.path.exists(out_path):

    os.makedirs(out_path)

files = os.listdir(path)

files.sort()

#files=files[:100]

for file in files:

    _, basename = os.path.split(file)

    if basename.lower().split('.')[-1] not in ['jpg', 'png']:

        continue

    stem, ext = os.path.splitext(basename)

    #stem=stem0.split('_')[2]

    gt_file = os.path.join(gt_path, stem+'.txt')

    img_path = os.path.join(path, file)

    print(img_path)

    #print(gt_file)

    img = cv.imread(img_path)

    if img is None:

        print('****************************')

        print('Image ' + img_path + ' may be a bad picture!')

        print('****************************')

        newname = os.path.join(path,stem+'.gif')

        os.rename(img_path,newname)

        img_path=newname

        print(img_path)

        print('Try read with imageio.')

        gif = imageio.mimread(img_path)

        if gif is None:

            print('****************************')

            print("Image " + img_path + " can't be read!")

            print('****************************')

        print('Read success!')

        img = cv.cvtColor(gif[0], cv.COLOR_RGB2BGR)

    img_size = img.shape

    im_size_min = np.min(img_size[0:2])

    im_size_max = np.max(img_size[0:2])

    im_scale = float(600) / float(im_size_min)

    if np.round(im_scale * im_size_max) > 1200:

        im_scale = float(1200) / float(im_size_max)

    re_im = cv.resize(img, None, None, fx=im_scale, fy=im_scale, interpolation=cv.INTER_LINEAR)

    re_size = re_im.shape

    cv.imwrite(os.path.join(out_path, stem) + '.jpg', re_im)

    with open(gt_file, 'r') as f:

        lines = f.readlines()

    for line in lines:

        splitted_line = line.strip().lower().split(',')

        pt_x = np.zeros((4, 1))

        pt_y = np.zeros((4, 1))

        pt_x[0, 0] = int(float(splitted_line[0]) / img_size[1] * re_size[1])

        pt_y[0, 0] = int(float(splitted_line[1]) / img_size[0] * re_size[0])

        pt_x[1, 0] = int(float(splitted_line[2]) / img_size[1] * re_size[1])

        pt_y[1, 0] = int(float(splitted_line[3]) / img_size[0] * re_size[0])

        pt_x[2, 0] = int(float(splitted_line[4]) / img_size[1] * re_size[1])

        pt_y[2, 0] = int(float(splitted_line[5]) / img_size[0] * re_size[0])

        pt_x[3, 0] = int(float(splitted_line[6]) / img_size[1] * re_size[1])

        pt_y[3, 0] = int(float(splitted_line[7]) / img_size[0] * re_size[0])

        ind_x = np.argsort(pt_x, axis=0)

        pt_x = pt_x[ind_x]

        pt_y = pt_y[ind_x]

        if pt_y[0] < pt_y[1]:

            pt1 = (pt_x[0], pt_y[0])

            pt3 = (pt_x[1], pt_y[1])

        else:

            pt1 = (pt_x[1], pt_y[1])

            pt3 = (pt_x[0], pt_y[0])

        if pt_y[2] < pt_y[3]:

            pt2 = (pt_x[2], pt_y[2])

            pt4 = (pt_x[3], pt_y[3])

        else:

            pt2 = (pt_x[3], pt_y[3])

            pt4 = (pt_x[2], pt_y[2])

        xmin = int(min(pt1[0], pt2[0]))

        ymin = int(min(pt1[1], pt2[1]))

        xmax = int(max(pt2[0], pt4[0]))

        ymax = int(max(pt3[1], pt4[1]))

        if xmin < 0:

            xmin = 0

        if xmax > re_size[1] - 1:

            xmax = re_size[1] - 1

        if ymin < 0:

            ymin = 0

        if ymax > re_size[0] - 1:

            ymax = re_size[0] - 1

        width = xmax - xmin

        height = ymax - ymin

        # reimplement

        step = 16.0

        x_left = []

        x_right = []

        x_left.append(xmin)

        x_left_start = int(math.ceil(xmin / 16.0) * 16.0)

        if x_left_start == xmin:

            x_left_start = xmin + 16

        for i in np.arange(x_left_start, xmax, 16):

            x_left.append(i)

        x_left = np.array(x_left)

        x_right.append(x_left_start - 1)

        for i in range(1, len(x_left) - 1):

            x_right.append(x_left[i] + 15)

        x_right.append(xmax)

        x_right = np.array(x_right)

        idx = np.where(x_left == x_right)

        x_left = np.delete(x_left, idx, axis=0)

        x_right = np.delete(x_right, idx, axis=0)

        if not os.path.exists('label_tmp'):

            os.makedirs('label_tmp')

        with open(os.path.join('label_tmp', stem) + '.txt', 'a') as f:

            #for i in range(len(x_left)):

            f.writelines("tianchi\t")

            f.writelines(str(int( pt_x[0, 0])))

            f.writelines("\t")

            f.writelines(str(int( pt_y[0, 0])))

            f.writelines("\t")

            f.writelines(str(int( pt_x[1, 0])))

            f.writelines("\t")

            f.writelines(str(int( pt_y[1, 0])))

            f.writelines("\t")

            f.writelines(str(int( pt_x[2, 0])))

            f.writelines("\t")

            f.writelines(str(int( pt_y[2, 0])))

            f.writelines("\t")

            f.writelines(str(int( pt_x[3, 0])))

            f.writelines("\t")

            f.writelines(str(int( pt_y[3, 0])))

            f.writelines("\n")

#ToVoc.py 上述执行完后直接运行这个脚本文件完美生成VOC文件

from xml.dom.minidom import Document

import cv2

import os

import glob

import shutil

import numpy as np

def generate_xml(name, lines, img_size, class_sets, doncateothers=True):

    doc = Document()

    def append_xml_node_attr(child, parent=None, text=None):

        ele = doc.createElement(child)

        if not text is None:

            text_node = doc.createTextNode(text)

            ele.appendChild(text_node)

        parent = doc if parent is None else parent

        parent.appendChild(ele)

        return ele

    img_name = name + '.jpg'

    # create header

    annotation = append_xml_node_attr('annotation')

    append_xml_node_attr('folder', parent=annotation, text='tianchi')

    append_xml_node_attr('filename', parent=annotation, text=img_name)

    source = append_xml_node_attr('source', parent=annotation)

    append_xml_node_attr('database', parent=source, text='coco_text_database')

    append_xml_node_attr('annotation', parent=source, text='tianchi')

    append_xml_node_attr('image', parent=source, text='tianchi')

    append_xml_node_attr('flickrid', parent=source, text='')

    owner = append_xml_node_attr('owner', parent=annotation)

    append_xml_node_attr('name', parent=owner, text='ms')

    size = append_xml_node_attr('size', annotation)

    append_xml_node_attr('width', size, str(img_size[1]))

    append_xml_node_attr('height', size, str(img_size[0]))

    append_xml_node_attr('depth', size, str(img_size[2]))

    append_xml_node_attr('segmented', parent=annotation, text='')

    # create objects

    objs = []

    for line in lines:

        splitted_line = line.strip().lower().split()

        cls = splitted_line[0].lower()

        if not doncateothers and cls not in class_sets:

            continue

        cls = 'dontcare' if cls not in class_sets else cls

        if cls == 'dontcare':

            continue

        obj = append_xml_node_attr('object', parent=annotation)

        occlusion = int(0)

        x1, y1, x2, y2 = int(float(splitted_line[1]) + 1), int(float(splitted_line[2]) + 1), \

                         int(float(splitted_line[3]) + 1), int(float(splitted_line[4]) + 1)

        x0,y0,x1,y1,x2,y2,x3,y3 = int(float(splitted_line[1])+1),int(float(splitted_line[2])+1),\

        int(float(splitted_line[3])+1),int(float(splitted_line[4])+1),int(float(splitted_line[5])+1),\

        int(float(splitted_line[6])+1),int(float(splitted_line[7])+1),int(float(splitted_line[8])+1)

        truncation = float(0)

        difficult = 1 if _is_hard(cls, truncation, occlusion, x1, y1, x2, y2) else 0

        truncted = 0 if truncation < 0.5 else 1

        append_xml_node_attr('name', parent=obj, text=cls)

        append_xml_node_attr('pose', parent=obj, text='none')

        append_xml_node_attr('truncated', parent=obj, text=str(truncted))

        append_xml_node_attr('difficult', parent=obj, text=str(int(difficult)))

        bb = append_xml_node_attr('bndbox', parent=obj)

        append_xml_node_attr('x0', parent=bb, text=str(int(x0)))

        append_xml_node_attr('y0', parent=bb, text=str(y0))

        append_xml_node_attr('x1', parent=bb, text=str(x1))

        append_xml_node_attr('y1', parent=bb, text=str(y1))

        append_xml_node_attr('x1', parent=bb, text=str(x2))

        append_xml_node_attr('y1', parent=bb, text=str(y2))

        append_xml_node_attr('x1', parent=bb, text=str(x3))

        append_xml_node_attr('y1', parent=bb, text=str(y3))

        o = {'class': cls, 'box': np.asarray([x0, y0,x1,y1, x2, y2,x3,y3], dtype=float), \

             'truncation': truncation, 'difficult': difficult, 'occlusion': occlusion}

        objs.append(o)

    return doc, objs

def _is_hard(cls, truncation, occlusion, x1, y1, x2, y2):

    hard = False

    if y2 - y1 < 25 and occlusion >= 2:

        hard = True

        return hard

    if occlusion >= 3:

        hard = True

        return hard

    if truncation > 0.8:

        hard = True

        return hard

    return hard

def build_voc_dirs(outdir):

    mkdir = lambda dir: os.makedirs(dir) if not os.path.exists(dir) else None

    mkdir(outdir)

    mkdir(os.path.join(outdir, 'Annotations'))

    mkdir(os.path.join(outdir, 'ImageSets'))

    mkdir(os.path.join(outdir, 'ImageSets', 'Layout'))

    mkdir(os.path.join(outdir, 'ImageSets', 'Main'))

    mkdir(os.path.join(outdir, 'ImageSets', 'Segmentation'))

    mkdir(os.path.join(outdir, 'JPEGImages'))

    mkdir(os.path.join(outdir, 'SegmentationClass'))

    mkdir(os.path.join(outdir, 'SegmentationObject'))

    return os.path.join(outdir, 'Annotations'), os.path.join(outdir, 'JPEGImages'), os.path.join(outdir, 'ImageSets',

                                                                                                 'Main')

if __name__ == '__main__':

    _outdir = 'TEXTVOC/VOC2007'

    _draw = bool(0)

    _dest_label_dir, _dest_img_dir, _dest_set_dir = build_voc_dirs(_outdir)

    _doncateothers = bool(1)

    for dset in ['train']:

        _labeldir = 'label_tmp'

        _imagedir = 're_image'

        class_sets = ('tianchi', 'dontcare')

        class_sets_dict = dict((k, i) for i, k in enumerate(class_sets))

        allclasses = {}

        fs = [open(os.path.join(_dest_set_dir, cls + '_' + dset + '.txt'), 'w') for cls in class_sets]

        ftrain = open(os.path.join(_dest_set_dir, dset + '.txt'), 'w')

        files = glob.glob(os.path.join(_labeldir, '*.txt'))

        files.sort()

        for file in files:

            path, basename = os.path.split(file)

            stem, ext = os.path.splitext(basename)

            with open(file, 'r') as f:

                lines = f.readlines()

            img_file = os.path.join(_imagedir, stem + '.jpg')

            print(img_file)

            img = cv2.imread(img_file)

            img_size = img.shape

            doc, objs = generate_xml(stem, lines, img_size, class_sets=class_sets, doncateothers=_doncateothers)

            cv2.imwrite(os.path.join(_dest_img_dir, stem + '.jpg'), img)

            xmlfile = os.path.join(_dest_label_dir, stem + '.xml')

            with open(xmlfile, 'w') as f:

                f.write(doc.toprettyxml(indent='    '))

            ftrain.writelines(stem + '\n')

            cls_in_image = set([o['class'] for o in objs])

            for obj in objs:

                cls = obj['class']

                allclasses[cls] = 0 \

                    if not cls in list(allclasses.keys()) else allclasses[cls] + 1

            for cls in cls_in_image:

                if cls in class_sets:

                    fs[class_sets_dict[cls]].writelines(stem + ' 1\n')

            for cls in class_sets:

                if cls not in cls_in_image:

                    fs[class_sets_dict[cls]].writelines(stem + ' -1\n')

        (f.close() for f in fs)

        ftrain.close()

        print('~~~~~~~~~~~~~~~~~~~')

        print(allclasses)

        print('~~~~~~~~~~~~~~~~~~~')

        shutil.copyfile(os.path.join(_dest_set_dir, 'train.txt'), os.path.join(_dest_set_dir, 'val.txt'))

        shutil.copyfile(os.path.join(_dest_set_dir, 'train.txt'), os.path.join(_dest_set_dir, 'trainval.txt'))

        for cls in class_sets:

            shutil.copyfile(os.path.join(_dest_set_dir, cls + '_train.txt'),

                            os.path.join(_dest_set_dir, cls + '_trainval.txt'))

            shutil.copyfile(os.path.join(_dest_set_dir, cls + '_train.txt'),

                            os.path.join(_dest_set_dir, cls + '_val.txt'))