-
Notifications
You must be signed in to change notification settings - Fork 0
/
DNetYoloV3Dataset.py
32 lines (26 loc) · 1.22 KB
/
DNetYoloV3Dataset.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
import os
import pathlib
def WriteDataSet(target_filename, imagefolder, abspath):
f = open(target_filename,"w+")
for filename in os.listdir(imagefolder):
if filename.endswith(".jpg") or filename.endswith(".png"):
f.write(os.path.join(abspath, imagefolder, filename) +"\n")
def WriteNamesFile(classes, filename):
f = open(filename,"w+")
for name in classes:
f.write(name + "\n")
def WriteDataFile(datasetname, classes, tsetname, vsetname, namesname):
f = open(datasetname + ".data","w+")
f.write("classes=" + str(classes) + "\n")
f.write("train=" + tsetname + "\n")
f.write("valid=" + vsetname + "\n")
f.write("names=" + namesname)
def WriteYoloV3Set(datasetname, train_images, val_images, classes):
abspath = pathlib.Path(__file__).parent.absolute()
trainset = os.path.join(abspath, datasetname, datasetname + "_train.txt")
valset = os.path.join(abspath, datasetname, datasetname + "_val.txt")
names = os.path.join(abspath, datasetname, datasetname + ".names")
WriteDataSet(trainset, train_images, abspath)
WriteDataSet(valset, val_images, abspath)
WriteNamesFile(classes, names)
WriteDataFile(datasetname, len(classes), trainset, valset, names)