object-detection/summer_processing.py at main · liekon/object-detection · 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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
import os
import time
import xml.dom.minidom


BATCH_SIZE = 500
DAC_CONTEST = '/home/xilinx/jupyter_notebooks/summer_school'
IMG_DIR = '/home/xilinx/jupyter_notebooks/summer_school/images'
OVERLAY_DIR = '/home/xilinx/jupyter_notebooks/summer_school/overlay'
RESULT = '/home/xilinx/jupyter_notebooks/summer_school/result'
TIME_DIR = '/home/xilinx/jupyter_notebooks/summer_school/result/time'
COORD_DIR = '/home/xilinx/jupyter_notebooks/summer_school/result/coordinate'
XML_PATH = '/home/xilinx/jupyter_notebooks/summer_school/result/xml'

def path_mkdir():
    if not os.path.isdir(DAC_CONTEST):
        os.mkdir(DAC_CONTEST)
    if not os.path.isdir(IMG_DIR):
        os.mkdir(IMG_DIR)
    if not os.path.isdir(OVERLAY_DIR):
        os.mkdir(OVERLAY_DIR)
    if not os.path.isdir(RESULT):
        os.mkdir(RESULT)
    if not os.path.isdir(TIME_DIR):
        os.mkdir(TIME_DIR)
    if not os.path.isdir(COORD_DIR):
        os.mkdir(COORD_DIR)
    if not os.path.isdir(XML_PATH):
        os.mkdir(XML_PATH)


# Get image name list
def get_image_names():
    names_temp = [f for f in os.listdir(IMG_DIR) if f.endswith('.jpg')]
    names_temp.sort(key= lambda x:int(x[:-4]))
    return names_temp


# Process the images in batches, may help when write to XML
def get_image_batch():
    image_list = get_image_names()
    batches = list()
    for i in range(0, len(image_list), BATCH_SIZE):
        batches.append(image_list[i:i+BATCH_SIZE])
    return batches


# Get image paths in batches
def get_image_path(image_name):
    return os.path.join(IMG_DIR, image_name)


# Return a batch of image dir  when `send` is called
class Agent:
    def __init__(self, teamname):
        self.batch_count = 0
        self.dac_contest = DAC_CONTEST
        self.img_dir = IMG_DIR
        self.overlay_dir = OVERLAY_DIR
        self.overlay_dir_team = OVERLAY_DIR + '/' + teamname
        self.result = RESULT
        self.time_dir = TIME_DIR
        self.coord_dir = COORD_DIR
        self.xml_path = XML_PATH
        self.coord_team = COORD_DIR + '/' + teamname
        self.xml_team = XML_PATH + '/' + teamname
        self.contestant = DAC_CONTEST + '/' + teamname
        folder_list = [self.dac_contest, self.img_dir, self.overlay_dir,
                       self.overlay_dir_team,
                       self.result,
                       self.time_dir, self.coord_dir, self.xml_path,
                       self.coord_team, self.xml_team, self.contestant]
        for folder in folder_list:
            if not os.path.isdir(folder):
                os.mkdir(folder)
        self.img_list = get_image_names()
        self.img_batch = get_image_batch()

    def send(self, interval_time, batches):
        time.sleep(interval_time)
        tmp = batches[self.batch_count]
        self.batch_count += 1
        return tmp

    def reset_batch_count(self):
        self.batch_count = 0

    def write(self, t_batch, total_img, teamname):
        fps = total_img / t_batch
        with open(self.time_dir + '/' + teamname + '.txt', 'a+') as f:
            f.write("\n" + teamname + " Frames per second: " +
                    str(fps) + '\n')

    def save_results_xml(self, result_rectangle):
        if len(result_rectangle) != len(self.img_list):
            raise ValueError("Result length not equal to number of images.")
        for i in range(len(self.img_list)):
            doc = xml.dom.minidom.Document()
            root = doc.createElement('annotation')

            doc.appendChild(root)
            name_e = doc.createElement('filename')
            name_t = doc.createTextNode(self.img_list[i])
            name_e.appendChild(name_t)
            root.appendChild(name_e)

            size_e = doc.createElement('size')
            node_width = doc.createElement('width')
            node_width.appendChild(doc.createTextNode("640"))
            node_length = doc.createElement('length')
            node_length.appendChild(doc.createTextNode("360"))
            size_e.appendChild(node_width)
            size_e.appendChild(node_length)
            root.appendChild(size_e)

            object_node = doc.createElement('object')
            node_name = doc.createElement('name')
            node_name.appendChild(doc.createTextNode("NotCare"))
            node_bnd_box = doc.createElement('bndbox')
            node_bnd_box_xmin = doc.createElement('xmin')
            node_bnd_box_xmin.appendChild(
                doc.createTextNode(str(result_rectangle[i][0])))
            node_bnd_box_xmax = doc.createElement('xmax')
            node_bnd_box_xmax.appendChild(
                doc.createTextNode(str(result_rectangle[i][1])))
            node_bnd_box_ymin = doc.createElement('ymin')
            node_bnd_box_ymin.appendChild(
                doc.createTextNode(str(result_rectangle[i][2])))
            node_bnd_box_ymax = doc.createElement('ymax')
            node_bnd_box_ymax.appendChild(
                doc.createTextNode(str(result_rectangle[i][3])))
            node_bnd_box.appendChild(node_bnd_box_xmin)
            node_bnd_box.appendChild(node_bnd_box_xmax)
            node_bnd_box.appendChild(node_bnd_box_ymin)
            node_bnd_box.appendChild(node_bnd_box_ymax)

            object_node.appendChild(node_name)
            object_node.appendChild(node_bnd_box)
            root.appendChild(object_node)

            file_name = self.img_list[i].replace('jpg', 'xml')
            with open(self.xml_team + "/" + file_name, 'w') as fp:
                doc.writexml(fp, indent='\t', addindent='\t',
                             newl='\n', encoding="utf-8")