From 2cb1840c6830cd5ff6be8bbebce1c457fbe00a8c Mon Sep 17 00:00:00 2001
From: Constantin Wenger <constantin.wenger@googlemail.com>
Date: Tue, 01 Feb 2022 22:46:42 +0000
Subject: [PATCH] updated requirements

---
 setup_train.py |   11 +++++++----
 1 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/setup_train.py b/setup_train.py
index d7897c9..27005c5 100644
--- a/setup_train.py
+++ b/setup_train.py
@@ -1,11 +1,14 @@
-import os
 from glob import glob
+import os
 import random
 
+from config import Config
+
+
 def main():
     random.seed()
     data_list = []
-    for subdir in glob('data/train/*'):
+    for subdir in glob('%s/train/*_update' % Config.data_dir):
         for data in glob(subdir + "/*.jpg"):
             data_list.append(os.path.abspath(data))
     random.shuffle(data_list)
@@ -13,10 +16,10 @@
     test_ratio = 0.1
     test_list = data_list[:int(test_ratio * len(data_list))]
     train_list = data_list[int(test_ratio * len(data_list)):]
-    with open('darknet/train.txt', 'w') as train_txt:
+    with open('%s/train.txt' % Config.darknet_dir, 'w') as train_txt:
         for data in train_list:
             train_txt.write(data + '\n')
-    with open('darknet/test.txt', 'w') as test_txt:
+    with open('%s/test.txt' % Config.darknet_dir, 'w') as test_txt:
         for data in test_list:
             test_txt.write(data + '\n')
     return

--
Gitblit v1.10.0