From 001193b9c7ba4036f1b0024b7801967c75694d9b Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Thu, 23 May 2019 13:15:44 +0200
Subject: [PATCH] updates

---
 train.py          |  6 ++----
 utils/datasets.py | 14 ++++++--------
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/train.py b/train.py
index 5a1e0cd4..6e8a884c 100644
--- a/train.py
+++ b/train.py
@@ -55,7 +55,6 @@ hyp = {'xy': 0.2,  # xy loss gain
 #        'momentum': 0.9025,  # SGD momentum
 #        'weight_decay': 0.0005417}  # optimizer weight decay
 
-
 def train(
         cfg,
         data_cfg,
@@ -143,7 +142,6 @@ def train(
                                   batch_size,
                                   augment=True,
                                   rect=False,
-                                  cache=True,
                                   multi_scale=multi_scale)
 
     # Initialize distributed training
@@ -168,7 +166,7 @@ def train(
         model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
 
     # Remove old results
-    for f in glob.glob('train_batch*.jpg') + glob.glob('test_batch*.jpg') + ['results.txt']:
+    for f in glob.glob('*_batch*.jpg') + glob.glob('results.txt'):
         os.remove(f)
 
     # Start training
@@ -307,7 +305,7 @@ if __name__ == '__main__':
     parser.add_argument('--batch-size', type=int, default=16, help='size of each image batch')
     parser.add_argument('--accumulate', type=int, default=1, help='accumulate gradient x batches before optimizing')
     parser.add_argument('--cfg', type=str, default='cfg/yolov3-spp.cfg', help='cfg file path')
-    parser.add_argument('--data-cfg', type=str, default='data/coco_32img.data', help='coco.data file path')
+    parser.add_argument('--data-cfg', type=str, default='data/coco.data', help='coco.data file path')
     parser.add_argument('--multi-scale', action='store_true', help='random image sizes per batch 320 - 608')
     parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)')
     parser.add_argument('--resume', action='store_true', help='resume training flag')
diff --git a/utils/datasets.py b/utils/datasets.py
index e66af194..49194207 100755
--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -130,7 +130,7 @@ class LoadWebcam:  # for inference
 
 
 class LoadImagesAndLabels(Dataset):  # for training/testing
-    def __init__(self, path, img_size=416, batch_size=16, augment=False, rect=True, image_weights=False, cache=False,
+    def __init__(self, path, img_size=416, batch_size=16, augment=False, rect=True, image_weights=False,
                  multi_scale=False):
         with open(path, 'r') as f:
             img_files = f.read().splitlines()
@@ -190,11 +190,8 @@ class LoadImagesAndLabels(Dataset):  # for training/testing
 
             self.batch_shapes = np.ceil(np.array(shapes) * img_size / 32.).astype(np.int) * 32
 
-        # Preload images
-        if cache and (n < 1001):  # preload all images into memory if possible
-            self.imgs = [cv2.imread(self.img_files[i]) for i in tqdm(range(n), desc='Reading images')]
-
         # Preload labels (required for weighted CE training)
+        self.imgs = [None] * n
         self.labels = [np.zeros((0, 5))] * n
         iter = tqdm(self.label_files, desc='Reading labels') if n > 1000 else self.label_files
         for i, file in enumerate(iter):
@@ -227,10 +224,11 @@ class LoadImagesAndLabels(Dataset):  # for training/testing
         label_path = self.label_files[index]
 
         # Load image
-        if hasattr(self, 'imgs'):  # preloaded
-            img = self.imgs[index]
-        else:
+        img = self.imgs[index]
+        if img is None:
             img = cv2.imread(img_path)  # BGR
+            if self.n < 1001:
+                self.imgs[index] = img  # cache image into memory
         assert img is not None, 'File Not Found ' + img_path
 
         # Augment colorspace