import os import piexif import logging import numpy as np from collections import OrderedDict from functools import partial from PIL import Image def check_readability(fpath): try: piexif.remove(fpath) except: pass try: im = Image.open(fpath) return 0 except Exception as e: os.remove(fpath) return 1 def is_accepted_image(fname, extensions): return any([fname.lower().endswith(ext) for ext in extensions]) def get_content(folder, extensions): _check = partial(is_accepted_image, extensions=extensions) content = [] for root, _, fnames in os.walk(folder): fnames = list(filter(_check, fnames)) if not fnames: continue content.append([root, fnames]) logging.info(f"Found {len(content)} sub-folders") return content def get_classnames(classnames): _names = np.loadtxt(classnames, dtype=str) class_names = OrderedDict() for name in _names: key = name.strip().partition(".")[-1].replace("_", " ").lower() class_names[key] = name logging.info(f"Parsed {len(class_names)} class names") return class_names