15 lines
593 B
Python
15 lines
593 B
Python
import os
|
|
from typing import Generator
|
|
|
|
SUPPORTED_EXT = [".jpg", ".png", ".jpeg", ".bmp", ".jfif", ".webp"]
|
|
|
|
def image_generator(image_dir: str, do_recurse: bool = True) -> Generator[str, None, None]:
|
|
if do_recurse:
|
|
for root, dirs, files in os.walk(image_dir):
|
|
for file in files:
|
|
if any(file.endswith(ext) for ext in SUPPORTED_EXT):
|
|
yield os.path.join(root, file)
|
|
else:
|
|
for file in os.listdir(image_dir):
|
|
if any(file.endswith(ext) for ext in SUPPORTED_EXT):
|
|
yield os.path.join(image_dir, file) |