File size: 1,015 Bytes
ab576ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
from pathlib import Path
from typing import Generator, Union
from PIL import Image
class ImageStreamer:
def __init__(self, image_or_folder: Union[str, Path]):
path = Path(image_or_folder)
self.generator = self._get_image_generator(path)
def _get_image_generator(self, path: Path) -> Generator[Image.Image, None, None]:
if path.is_file():
image_paths = [path] if self._is_image_file(path) else []
elif path.is_dir():
image_paths = [
p
for p in path.rglob('**/*')
if self._is_image_file(p)
]
else:
raise TypeError(f'Invalid path to images {path}')
for p in image_paths:
yield Image.open(p)
def _is_image_file(self, path: Path) -> bool:
try:
image = Image.open(path)
image.verify()
return True
except Exception:
return False
def __iter__(self):
return self.generator
|