picture_ratio.py

#!/usr/bin/env python3

import sys

from glob import glob
from random import sample

import numpy as np
import tensorflow as tf

from PIL import Image
from sklearn.model_selection import train_test_split


SIZE = 224
RATIO = 1.2
COUNT = 450


def get_image(filename):
    return Image.open(filename)

def get_x(image):
    resized = image.resize((SIZE, SIZE), Image.ANTIALIAS).convert('L')
    return np.array(resized) / 128.0 - 1

def get_y(image):
    width, height = image.size
    if width > height * RATIO:
        return 0
    elif height > width * RATIO:
        return 2
    else:
        return 1


GLOB_PATTERN = "/home/aragaer/Pictures/wedding/*.jpg"

def load_data():
    """Returns the picture dataset as (train_x, train_y), (test_x, test_y)."""
    filenames = sample(glob(GLOB_PATTERN, recursive=True), COUNT)
    x = np.ndarray((COUNT, SIZE, SIZE), dtype=np.float32)
    y = np.ndarray((COUNT, 1), dtype=np.int32)
    for i, filename in enumerate(filenames):
        image = get_image(filename)
        x[i] = get_x(image)
        y[i] = get_y(image)
        if i % 100 == 0:
            print("Loaded", i)

    train_x, test_x, train_y, test_y = train_test_split(x, y, test_size=0.2)

    return (train_x, train_y), (test_x, test_y)

def train_input_fn(pixels, labels, batch_size):
    """An input function for training"""
    pixels = tf.convert_to_tensor(np.asarray(pixels))
    labels = tf.convert_to_tensor(np.asarray(labels))
    dataset = tf.data.Dataset.from_tensor_slices(({"Pixels": pixels}, labels))
    return dataset.shuffle(1000).repeat().batch(batch_size)

def test_input_fn(pixels, labels, batch_size):
    """An input function for training"""
    pixels = tf.convert_to_tensor(np.asarray(pixels))
    labels = tf.convert_to_tensor(np.asarray(labels))
    dataset = tf.data.Dataset.from_tensor_slices(({"Pixels": pixels}, labels))
    return dataset.batch(batch_size)

def predict_input_fn(pixels, batch_size):
    pixels = tf.convert_to_tensor(np.asarray(pixels))
    dataset = tf.data.Dataset.from_tensor_slices({"Pixels": pixels})
    return dataset.batch(batch_size)


def my_model_fn(features, labels, mode, params):
    net = tf.feature_column.input_layer(features, params['feature_columns'])
    net = tf.reshape(net, [-1, SIZE, SIZE, 1])
    net = tf.layers.conv2d(net, filters=params['conv']['filters'],
                           kernel_size=params['conv']['kernel'])
    net = tf.layers.max_pooling2d(inputs=net, pool_size=4, strides=4)
    count = 1
    for s in net.shape[1:]:
        count *= int(s)
    
    net = tf.reshape(net, [-1, count])
    net = tf.layers.dense(net, units=params['dense'], activation=tf.nn.relu)
    logits = tf.layers.dense(net, params['n_classes'], activation=None)
    # Compute predictions.
    predicted_classes = tf.argmax(logits, 1)
    if mode == tf.estimator.ModeKeys.PREDICT:
        predictions = {
            'class_ids': predicted_classes[:, tf.newaxis],
            'probabilities': tf.nn.softmax(logits),
            'logits': logits,
        }
        return tf.estimator.EstimatorSpec(mode,
                                          predictions=predictions)

    # Compute loss.
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels,
                                                  logits=logits)

    # Compute evaluation metrics.
    accuracy = tf.metrics.accuracy(labels=labels,
                                   predictions=predicted_classes,
                                   name='acc_op')
    metrics = {'accuracy': accuracy}
    tf.summary.scalar('accuracy', accuracy[1])

    if mode == tf.estimator.ModeKeys.EVAL:
        return tf.estimator.EstimatorSpec(mode,
                                          loss=loss,
                                          eval_metric_ops=metrics)

    optimizer = tf.train.AdagradOptimizer(learning_rate=0.1)
    train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
    return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)


def main():
    picture_feature_column = tf.feature_column.numeric_column(key="Pixels",
                                                              shape=[SIZE, SIZE])
    train, test = load_data()
    features, labels = train

    batch_size = 100

    classifier = tf.estimator.Estimator(
        model_fn=my_model_fn,
        params={
            'feature_columns': [picture_feature_column],
            'dense': 10,
            'conv': {'filters': 1, 'kernel': 8},
            'n_classes': 3,
        }, model_dir="beauty")

    for i in range(100):
        print(i, "Training...", end='')
        sys.stdout.flush()
        classifier.train(input_fn=lambda: train_input_fn(features, labels, batch_size),
                         steps=1000)
        print("Done")
        print(i, "Evaluating...", end='')
        sys.stdout.flush()
        test_result = classifier.evaluate(input_fn=lambda: test_input_fn(*test, batch_size))
        print("Done")
        print('Test set accuracy: {accuracy:0.3f}\n'.format(**test_result))

    pixels = np.ndarray((1, SIZE, SIZE), dtype=np.float32)
    correct = 0
    total = 0
    for filename in glob(GLOB_PATTERN, recursive=True)[:1000]:
        image = get_image(filename)
        pixels[0] = get_x(image)
        actual = get_y(image)
        predicted = classifier.predict(input_fn=lambda: predict_input_fn(pixels, batch_size))
        result = next(predicted)
        class_id = result['class_ids'][0]
        correct += actual == class_id
        total += 1
        print(actual, class_id, result['probabilities'][class_id]*100)

    print(correct,"/",total)


if __name__ == '__main__':
    main()