show error when dataset directory is empty

Former-commit-id: 191fb3cbc9486e364faa51afbe4d060fd2df7348
javiribera · May 10, 2018 · 376e145 · 376e145
1 parent 38a88b2
commit 376e145
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 12 deletions.
diff --git a/object-locator/data.py b/object-locator/data.py
@@ -300,16 +300,20 @@ def __init__(self,
                 xml_filename = filename
                 break
 
+        # Ignore files that are not images
+        listfiles = [f for f in listfiles
+                     if any(f.lower().endswith(ext) for ext in IMG_EXTENSIONS)]
+
+        if len(listfiles) == 0:
+            raise ValueError(f"There are no images in '{directory}'")
+
         if xml_filename is None:
             print('W: The dataset directory %s does not contain '
                   'a XML file with groundtruth. Metrics will not be evaluated.'
                   'Only estimations will be returned.' % directory)
 
         self.there_is_gt = (xml_filename is not None) and (not ignore_gt)
 
-        # Ignore files that are not images
-        listfiles = [f for f in listfiles
-                     if any(f.lower().endswith(ext) for ext in IMG_EXTENSIONS)]
 
         # XML does not exist (no GT available)
         if not self.there_is_gt:

diff --git a/object-locator/locate.py b/object-locator/locate.py
@@ -56,15 +56,19 @@
     os.makedirs(os.path.join(args.out_dir, 'painted'), exist_ok=True)
 
 # Data loading code
-testset = XMLDataset(args.dataset,
-                     transforms=transforms.Compose([
-                         ScaleImageAndLabel(size=(args.height, args.width)),
-                         transforms.ToTensor(),
-                         transforms.Normalize((0.5, 0.5, 0.5),
-                                              (0.5, 0.5, 0.5)),
-                     ]),
-                     ignore_gt=not args.evaluate,
-                     max_dataset_size=args.max_testset_size)
+try:
+    testset = XMLDataset(args.dataset,
+                         transforms=transforms.Compose([
+                             ScaleImageAndLabel(size=(args.height, args.width)),
+                             transforms.ToTensor(),
+                             transforms.Normalize((0.5, 0.5, 0.5),
+                                                  (0.5, 0.5, 0.5)),
+                         ]),
+                         ignore_gt=not args.evaluate,
+                         max_dataset_size=args.max_testset_size)
+except ValueError as e:
+    print(f'E: {e}')
+    exit(-1)
 testset_loader = data.DataLoader(testset,
                                  batch_size=1,
                                  num_workers=args.nThreads,