From 54a9e4f8764aae695e42d62869563f2a9bbe8ec4 Mon Sep 17 00:00:00 2001
From: lorenzomammana <lorenzom96@hotmail.it>
Date: Thu, 9 Jul 2020 19:39:28 +0200
Subject: [PATCH] Refactor code to reduce duplication

---
 utils/datasets.py | 33 ++++++++++++++++-----------------
 1 file changed, 16 insertions(+), 17 deletions(-)

diff --git a/utils/datasets.py b/utils/datasets.py
index d96ae41..d3bc574 100755
--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -280,32 +280,31 @@ class LoadImagesAndLabels(Dataset):  # for training/testing
     def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
                  cache_images=False, single_cls=False, stride=32, pad=0.0):
         try:
-            if type(path) is list:
-                # Multiple datasets handler
-                f = []
-                for subpath in path:
+            f = []
+            for subpath in path if isinstance(path, list) else [path]:
+                subpath = str(Path(subpath))  # os-agnostic
+                parent = str(Path(subpath).parent) + os.sep
+                if os.path.isfile(subpath):  # file
                     with open(subpath, 'r') as t:
-                        subpath = str(Path(subpath))  # os-agnostic
-                        parent = str(Path(subpath).parent) + os.sep
                         t = t.read().splitlines()
                         t = [x.replace('./', parent) if x.startswith('./') else x for x in t]  # local to global path
                         f += t
-                path = str(Path(path[0]))  # from now on treat multiple datasets as single
-            else:
-                path = str(Path(path))  # os-agnostic
-                parent = str(Path(path).parent) + os.sep
-                if os.path.isfile(path):  # file
-                    with open(path, 'r') as f:
-                        f = f.read().splitlines()
-                        f = [x.replace('./', parent) if x.startswith('./') else x for x in f]  # local to global path
-                elif os.path.isdir(path):  # folder
-                    f = glob.iglob(path + os.sep + '*.*')
+                elif os.path.isdir(subpath):  # folder
+                    f = glob.iglob(subpath + os.sep + '*.*')
+                    # Maybe change this to f += glob.glob, this should allow handling also multiple folders
                 else:
-                    raise Exception('%s does not exist' % path)
+                    raise Exception('%s does not exist' % subpath)
             self.img_files = [x.replace('/', os.sep) for x in f if os.path.splitext(x)[-1].lower() in img_formats]
         except:
+            # Maybe avoid handling bare exceptions
             raise Exception('Error loading data from %s. See %s' % (path, help_url))
 
+        # Still need to do this for compatibility with the .npy and shape file saves
+        if isinstance(path, list):
+            path = str(Path(path[0]))
+        else:
+            path = str(Path(path))
+
         n = len(self.img_files)
         assert n > 0, 'No images found in %s. See %s' % (path, help_url)
         bi = np.floor(np.arange(n) / batch_size).astype(np.int)  # batch index