From b7c9eac280c9ab03e797a3bd4ae26d5cff17563b Mon Sep 17 00:00:00 2001 From: andreasjansson Date: Wed, 20 Oct 2021 00:24:47 +0200 Subject: [PATCH] Continue on error in dataset_tool.py When trying to use dataset_tool.py on my own folder of images, the script died half way through with an error. This patch makes dataset_tool.py more robust to errors, and will write warning messages instead of failing if an error is encountered. --- dataset_tool.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/dataset_tool.py b/dataset_tool.py index 747189fd..c3bd7d72 100644 --- a/dataset_tool.py +++ b/dataset_tool.py @@ -86,8 +86,12 @@ def iterate_images(): for idx, fname in enumerate(input_images): arch_fname = os.path.relpath(fname, source_dir) arch_fname = arch_fname.replace('\\', '/') - img = np.array(PIL.Image.open(fname)) - yield dict(img=img, label=labels.get(arch_fname)) + try: + img = np.array(PIL.Image.open(fname)) + except Exception as e: + sys.stderr.write(f'Failed to read {fname}: {e}\n') + continue + yield dict(img=img, label=labels.get(arch_fname), fname=fname) if idx >= max_idx-1: break return max_idx, iterate_images() @@ -409,7 +413,15 @@ def convert_dataset( archive_fname = f'{idx_str[:5]}/img{idx_str}.png' # Apply crop and resize. - img = transform_image(image['img']) + try: + img = transform_image(image['img']) + except Exception as e: + if "fname" in image: + error_msg = f'Failed to process image {image["fname"]}: {e}' + else: + error_msg = f'Failed to process image at index {idx}: {e}' + sys.stderr.write(error_msg + '\n') + continue # Transform may drop images. if img is None: