yolo-pet/pet_label.py at master · ideaRunner/yolo-pet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import shutil
import random
classes0 = ['Abyssinian', 'Bengal', 'Birman', 'Bombay', 'British', 'Egyptian', 'Maine', 'Persian', 'Ragdoll', 'Russian', 'Siamese', 'Sphynx', 'american', 'basset', 'beagle', 'boxer', 'chihuahua', 'english', 'german', 'great', 'havanese', 'japanese', 'keeshond', 'leonberger', 'miniature', 'newfoundland', 'pomeranian', 'pug', 'saint', 'samoyed', 'scottish', 'shiba', 'staffordshire', 'wheaten', 'yorkshire']
classes = ['Abyssinian', 'Bengal', 'Birman', 'Bombay', 'British_Shorthair', 'Egyptian_Mau', 'Maine_Coon', 'Persian', 'Ragdoll', 'Russian_Blue', 'Siamese', 'Sphynx', 'american_bulldog', 'american_pit_bull_terrier', 'basset_hound', 'beagle', 'boxer', 'chihuahua', 'english_cocker_spaniel', 'english_setter', 'german_shorthaired', 'great_pyrenees', 'havanese', 'japanese_chin', 'keeshond', 'leonberger', 'miniature_pinscher', 'newfoundland', 'pomeranian', 'pug', 'saint_bernard', 'samoyed', 'scottish_terrier', 'shiba_inu', 'staffordshire_bull_terrier', 'wheaten_terrier', 'yorkshire_terrier']

def convert(size, box):
    dw = 1./size[0]
    dh = 1./size[1]
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    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)


def convert_annotation(image_id):
    try:
        in_file = open('annotations/xmls/%s.xml' % image_id)
        out_file = open('labels/%s.txt' % image_id, '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)
        breeds = root.find('filename').text
        cls = breeds[:breeds.rfind('_')]

        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            # breeds = string.capwords(obj.find('name').text)+'_'+cls
            if cls not in classes or int(difficult) == 1:
                print("can not find"+str(cls))
                continue
            # if int(difficult) == 1:
            #     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))
            bb = convert((w,h), b)
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
            return True
    except:
        return False


def generate():
    wd = getcwd()
    if not os.path.exists('labels'):
        os.makedirs('labels')

    train_image_ids = []
    with open('annotations/trainval.txt') as f:
        lines = f.readlines()
        for x in lines:
            train_image_ids.append(x.split(' ')[0])

    test_image_ids = []
    with open('annotations/test.txt') as f:
        lines = f.readlines()
        for x in lines:
            test_image_ids.append(x.split(' ')[0])

    train_list_file = open('Train_List.txt', 'w')
    val_list_file = open('Val_List.txt', 'w')
    num_train =0
    num_val =0
    for image_id in train_image_ids:
        if convert_annotation(image_id):
            if random.random() > 0.1:
                train_list_file.write('%s/JPEGImages/%s.jpg\n' % (wd, image_id))
                num_train = num_train + 1
            else:
                val_list_file.write('%s/JPEGImages/%s.jpg\n' % (wd, image_id))
                num_val = num_val + 1
    print("Train files: {0}, Val files: {1}.".format(num_train, num_val))
    train_list_file.close()
    val_list_file.close()


def clean():
    wd = getcwd()
    os.remove(wd+'/Train_List.txt')
    os.remove(wd+'/Val_List.txt')
    shutil.rmtree('labels')
    print("Clean Finished.")


if __name__ == '__main__':
    import argparse
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="'generate' labels or 'clean' on PetSet")
    parser.add_argument("command",
                        metavar="<command>",
                        help="'generate' labels or 'clean' on PetSet")

    args = parser.parse_args()
    print("Command: ", args.command)
    if args.command == "generate":
        generate()
    # elif args.command == "back":
    #     val_back_train()
    elif args.command == "clean":
        clean()
    elif args.command == "rg":
        clean()
        generate()
    else:
        print("Command error, 'generate' or 'clean'")


# cannot find some xml files