some updates

2019-05-18 18:49:40 +08:00
parent ec5c4ac97c
commit f136b8f60a
4 changed files with 28 additions and 7 deletions
--- a/tools/TrainCNN/generate.py
+++ b/tools/TrainCNN/generate.py
@@ -41,6 +41,13 @@ class DataSet:

    def generate_data_sets(self, folder):
        sets = []
+        mini = 1000000
+        for i in range(OUTPUT_NODES):
+            dir = "%s/%d" % (folder, i)
+            files = os.listdir(dir)
+            if mini > len(files):
+                mini = len(files)
+
        for i in range(OUTPUT_NODES):
            dir = "%s/%d" % (folder, i)
            files = os.listdir(dir)