From e6ef7bc73d5790c864d3977f7cc4e2e9905c03f5 Mon Sep 17 00:00:00 2001 From: Marcus Bloice Date: Thu, 7 Mar 2019 12:28:02 +0100 Subject: [PATCH] Fixes #156 where the Pipeline status function would error on unlabelled datasets. --- Augmentor/Pipeline.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/Augmentor/Pipeline.py b/Augmentor/Pipeline.py index 8a75740..1ae2314 100644 --- a/Augmentor/Pipeline.py +++ b/Augmentor/Pipeline.py @@ -800,12 +800,21 @@ def status(self): print("Images: %s" % len(self.augmentor_images)) - label_pairs = sorted(set([x.label_pair for x in self.augmentor_images])) + # TODO: find a better way that doesn't need to iterate over every image + # TODO: get rid of this label_pair property as nowhere else uses it + # Check if we have any labels before printing label information. + label_count = 0 + for image in self.augmentor_images: + if image.label_pair is not None: + label_count += 1 - print("Classes: %s" % len(label_pairs)) + if len(label_count) != 0: + label_pairs = sorted(set([x.label_pair for x in self.augmentor_images])) - for label_pair in label_pairs: - print ("\tClass index: %s Class label: %s " % (label_pair[0], label_pair[1])) + print("Classes: %s" % len(label_pairs)) + + for label_pair in label_pairs: + print ("\tClass index: %s Class label: %s " % (label_pair[0], label_pair[1])) if len(self.augmentor_images) != 0: print("Dimensions: %s" % len(self.distinct_dimensions))