From 2b213fbf66a96033dcb95dfb422a8268625d2337 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Sun, 14 Feb 2021 00:05:45 -0500 Subject: [PATCH] added .csv image loading utils --- flash/vision/classification/data.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/flash/vision/classification/data.py b/flash/vision/classification/data.py index 6f90f2571d..766da2460b 100644 --- a/flash/vision/classification/data.py +++ b/flash/vision/classification/data.py @@ -17,7 +17,7 @@ import pandas as pd import torch -from PIL import Image +from PIL import Image, UnidentifiedImageError from pytorch_lightning.utilities.exceptions import MisconfigurationException from torchvision import transforms as T from torchvision.datasets import VisionDataset @@ -241,9 +241,13 @@ def before_collate(self, samples: Any) -> Any: if isinstance(samples, (list, tuple)) and all(isinstance(p, str) for p in samples): outputs = [] for sample in samples: - output = self._loader(sample) - transform = self._valid_transform if self._use_valid_transform else self._train_transform - outputs.append(transform(output)) + try: + output = self._loader(sample) + transform = self._valid_transform if self._use_valid_transform else self._train_transform + outputs.append(transform(output)) + except UnidentifiedImageError as e: + print(f'Skipping: could not read file {sample}') + return outputs raise MisconfigurationException("The samples should either be a tensor or a list of paths.")