utils.py

from functools import reduce
from PIL import Image, ImageStat
import tensorflow as tf
import cv2
from matplotlib import pyplot as plt
import numpy as np


def detect_color_image_v1(file, thumb_size=40, MSE_cutoff=22, adjust_color_bias=True):
    pil_img = Image.open(file)
    bands = pil_img.getbands()
    if bands == ('R', 'G', 'B') or bands == ('R', 'G', 'B', 'A'):
        thumb = pil_img.resize((thumb_size,thumb_size))
        SSE, bias = 0, [0, 0, 0]
        if adjust_color_bias:
            bias = ImageStat.Stat(thumb).mean[:3]
            bias = [b - sum(bias)/3 for b in bias]
        for pixel in thumb.getdata():
            mu = sum(pixel)/3
            SSE += sum((pixel[i] - mu - bias[i])*(pixel[i] - mu - bias[i]) for i in [0, 1, 2])
        MSE = float(SSE)/(thumb_size*thumb_size)
        if MSE <= MSE_cutoff:
            return("grayscale"),
        else:
            return("Color"),
    elif len(bands)==1:
        return("Black and white")
    else:
        return("Don't know...")


MONOCHROMATIC_MAX_VARIANCE = 0.005
COLOR = 1000
MAYBE_COLOR = 100

def detect_color_image(file):
    v = ImageStat.Stat(Image.open(file)).var
    is_monochromatic = reduce(lambda x, y: x and y < MONOCHROMATIC_MAX_VARIANCE, v, True)
    if is_monochromatic:
        return("Monochromatic image")
    else:
        if len(v) == 3:
            maxmin = abs(max(v) - min(v))
            if maxmin > COLOR:
                return("Color")
            elif maxmin > MAYBE_COLOR:
                return("Maybe color")
            else:
                return("grayscale")
        elif len(v) == 1:
            return("Black and white")
        else:
            return("Don't know...")


# exaple of data augmentation

img = cv2.imread('./images/potato.jpg')

# Original Image

plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
plt.show()

# Augmented Image
aug = tf.image.random_saturation(
    img,
    lower=0.1,
    upper=1)

plt.imshow(cv2.cvtColor(aug.numpy(), cv2.COLOR_BGR2RGB))
plt.show()


# Pipeline

def aug_flow(img):
    data = []
    for i in range(9):
        img = tf.image.random_brightness(img, max_delta=0.02)
        img = tf.image.random_contrast(img, lower=0.6, upper=1)
        img = tf.image.random_flip_left_right(img)
        img = tf.image.random_jpeg_quality(img, min_jpeg_quality=90, max_jpeg_quality=100)
        img = tf.image.random_saturation(img, lower=0.9, upper=1)

        data.append(img)
    return data


aug_images = aug_flow(img)

for idx, aug_img in enumerate(aug_images):
    plt.subplot(3, 3, idx+1)
    bgr = cv2.cvtColor(aug_img.numpy(), cv2.COLOR_BGR2RGB)
    plt.imshow(bgr)