forked from ferrygun/PDFTableExtract
-
Notifications
You must be signed in to change notification settings - Fork 1
/
build_logos.py
91 lines (75 loc) · 2.63 KB
/
build_logos.py
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
# USAGE
# python build_logos.py
# import the necessary packages
from config import logos_config as config
from bs4 import BeautifulSoup
import os
# initialize the set of classes we have encountered so far
CLASSES = set()
# create the list of datasets to build
datasets = [
("train", config.TRAIN_TXT, config.TRAIN_CSV),
("test", config.TEST_TXT, config.TEST_CSV),
]
# loop over the datasets
for (dType, inputTxt, outputCSV) in datasets:
# load the contents of the input data split
print("[INFO] starting '{}' set...".format(dType))
imageIDs = open(inputTxt).read().strip().split("\n")
print("[INFO] {} total images in '{}' set".format(
len(imageIDs), dType))
# open the output CSV file
csv = open(outputCSV, "w")
# loop over the image IDs
for imageID in imageIDs:
# build the path to the image path and annotation file
imagePath = "{}.jpg".format(os.path.sep.join([
config.IMAGES_PATH, imageID]))
annotPath = "{}.xml".format(os.path.sep.join([
config.ANNOT_PATH, imageID]))
# load the annotation file, build the soup, and initialize
# the set of coordinates for this particular image
contents = open(annotPath).read()
soup = BeautifulSoup(contents, "html.parser")
coords = set()
# extract the image dimensions
w = int(soup.find("width").string)
h = int(soup.find("height").string)
# loop over all `object` elements
for o in soup.find_all("object"):
# extract the label and bounding box coordinates
label = o.find("name").string
xMin = int(o.find("xmin").string)
yMin = int(o.find("ymin").string)
xMax = int(o.find("xmax").string)
yMax = int(o.find("ymax").string)
# truncate any bounding box coordinates that may fall
# outside the boundaries of the image
xMin = max(0, xMin)
yMin = max(0, yMin)
xMax = min(w, xMax)
yMax = min(h, yMax)
# build a (hashable) tuple from the coordinates
coord = (xMin, yMin, xMax, yMax)
# if the coordinates already exist in our `coords` set,
# ignore the annotation (this is a peculiarity of the
# logos dataset)
if coord in coords:
continue
# write the image path, bounding box coordinates, and
# label to the output CSV file
row = [os.path.abspath(imagePath), str(xMin), str(yMin),
str(xMax), str(yMax), label]
csv.write("{}\n".format(",".join(row)))
# add the bounding box coordinates to our set and update
# the set of class labels
coords.add(coord)
CLASSES.add(label)
# close the output CSV file
csv.close()
# write the classes to file
print("[INFO] writing classes...")
csv = open(config.CLASSES_CSV, "w")
rows = [",".join([c, str(i)]) for (i, c) in enumerate(CLASSES)]
csv.write("\n".join(rows))
csv.close()