image_crawl/utils/imaging.py

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