CofeehousePy/services/nsfw_detection/coffeehouse_nsfw/image_utils.py

142 lines
5.3 KiB
Python

import os
import cv2
import pydload
import logging
import numpy as np
from PIL import Image as pil_image
if pil_image is not None:
_PIL_INTERPOLATION_METHODS = {
"nearest": pil_image.NEAREST,
"bilinear": pil_image.BILINEAR,
"bicubic": pil_image.BICUBIC,
}
# These methods were only introduced in version 3.4.0 (2016).
if hasattr(pil_image, "HAMMING"):
_PIL_INTERPOLATION_METHODS["hamming"] = pil_image.HAMMING
if hasattr(pil_image, "BOX"):
_PIL_INTERPOLATION_METHODS["box"] = pil_image.BOX
# This method is new in version 1.1.3 (2013).
if hasattr(pil_image, "LANCZOS"):
_PIL_INTERPOLATION_METHODS["lanczos"] = pil_image.LANCZOS
def load_img(
path, grayscale=False, color_mode="rgb", target_size=None, interpolation="nearest"
):
"""Loads an image into PIL format.
:param path: Path to image file.
:param grayscale: DEPRECATED use `color_mode="grayscale"`.
:param color_mode: One of "grayscale", "rgb", "rgba". Default: "rgb".
The desired image format.
:param target_size: Either `None` (default to original size)
or tuple of ints `(img_height, img_width)`.
:param interpolation: Interpolation method used to resample the image if the
target size is different from that of the loaded image.
Supported methods are "nearest", "bilinear", and "bicubic".
If PIL version 1.1.3 or newer is installed, "lanczos" is also
supported. If PIL version 3.4.0 or newer is installed, "box" and
"hamming" are also supported. By default, "nearest" is used.
:return: A PIL Image instance.
"""
if grayscale is True:
logging.warn("grayscale is deprecated. Please use " 'color_mode = "grayscale"')
color_mode = "grayscale"
if pil_image is None:
raise ImportError(
"Could not import PIL.Image. " "The use of `load_img` requires PIL."
)
if isinstance(path, type("")):
img = pil_image.open(path)
else:
path = cv2.cvtColor(path, cv2.COLOR_BGR2RGB)
img = pil_image.fromarray(path)
if color_mode == "grayscale":
if img.mode != "L":
img = img.convert("L")
elif color_mode == "rgba":
if img.mode != "RGBA":
img = img.convert("RGBA")
elif color_mode == "rgb":
if img.mode != "RGB":
img = img.convert("RGB")
else:
raise ValueError('color_mode must be "grayscale", "rgb", or "rgba"')
if target_size is not None:
width_height_tuple = (target_size[1], target_size[0])
if img.size != width_height_tuple:
if interpolation not in _PIL_INTERPOLATION_METHODS:
raise ValueError(
"Invalid interpolation method {} specified. Supported "
"methods are {}".format(
interpolation, ", ".join(_PIL_INTERPOLATION_METHODS.keys())
)
)
resample = _PIL_INTERPOLATION_METHODS[interpolation]
img = img.resize(width_height_tuple, resample)
return img
def img_to_array(img, data_format="channels_last", dtype="float32"):
"""Converts a PIL Image instance to a Numpy array.
# Arguments
img: PIL Image instance.
data_format: Image data format,
either "channels_first" or "channels_last".
dtype: Dtype to use for the returned array.
# Returns
A 3D Numpy array.
# Raises
ValueError: if invalid `img` or `data_format` is passed.
"""
if data_format not in {"channels_first", "channels_last"}:
raise ValueError("Unknown data_format: %s" % data_format)
# Numpy array x has format (height, width, channel)
# or (channel, height, width)
# but original PIL image has format (width, height, channel)
x = np.asarray(img, dtype=dtype)
if len(x.shape) == 3:
if data_format == "channels_first":
x = x.transpose(2, 0, 1)
elif len(x.shape) == 2:
if data_format == "channels_first":
x = x.reshape((1, x.shape[0], x.shape[1]))
else:
x = x.reshape((x.shape[0], x.shape[1], 1))
else:
raise ValueError("Unsupported image shape: %s" % (x.shape,))
return x
def load_images(image_paths, image_size, image_names):
"""
Function for loading images into numpy arrays for passing to model.predict
inputs:
image_paths: list of image paths to load
image_size: size into which images should be resized
outputs:
loaded_images: loaded images on which keras model can run predictions
loaded_image_indexes: paths of images which the function is able to process
"""
loaded_images = []
loaded_image_paths = []
for i, img_path in enumerate(image_paths):
try:
image = load_img(img_path, target_size=image_size)
image = img_to_array(image)
image /= 255
loaded_images.append(image)
loaded_image_paths.append(image_names[i])
except Exception as ex:
logging.exception(f"Error reading {img_path} {ex}", exc_info=True)
return np.asarray(loaded_images), loaded_image_paths