preprocess_cnn.py

# coding: utf-8

import os
import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
import matplotlib.image as mpimg

get_ipython().magic('matplotlib inline')

from PIL import Image
from skimage.transform import rescale, resize, rotate
from skimage.color import gray2rgb, rgb2gray

import time

from sklearn import preprocessing
import tensorflow as tf

IMAGE_SIZE = 138


class TFRecordGenerator:
    def __init__(self, image_size = IMAGE_SIZE):
        df_train = pd.read_csv('train.csv')
        df_test  = pd.read_csv('test.csv')
        df_validation = df_train.groupby(['species']).first()

        train_ids = df_train['id'].values
        labels      = df_train['species'].values
        
        le = preprocessing.LabelEncoder()
        le.fit(labels)
                
        train_size = train_ids.shape[0]
        train_matrix = np.empty([train_size, 2], np.uint32)
        train_matrix[:, 0] = train_ids
        train_matrix[:, 1] = le.transform(labels)
        
        self.image_size = image_size
        self.le = le
        self.train_matrix = train_matrix
        self.test_matrix  = df_test['id'].values
            
    def raw_image_path(self, id):
         return os.path.join("./images/",  str(id) + ".jpg")
        
    def load_image(self, id):        
        path = self.raw_image_path(id)
        image_2d = np.array(Image.open(path))
        image_3d = gray2rgb(image_2d)

        return image_3d # np.array
        
    def fit_image(self, image):
        fit_size = self.image_size
        
        # rescale image
        max_size = np.maximum(image.shape[0], image.shape[1])
        scale = fit_size / max_size
        image_3d = rescale(image, scale, mode='reflect') 
        
        # fit
        margin = np.array((fit_size, fit_size)) - image_3d.shape[0:2]
        margin = np.round(margin / 2).astype(int)

        pos_x = (margin[0], margin[0] + image_3d.shape[0])
        pos_y = (margin[1], margin[1] + image_3d.shape[1])
        
        image_norm = np.zeros((fit_size, fit_size, 3), )    
        image_norm[pos_x[0]:pos_x[1], pos_y[0]:pos_y[1], :] = image_3d
        
        return image_norm.astype(np.int32)
    
    def plot_image_raw(self, id):
        file_name = self.raw_image_path(id)
        print(file_name)
        
        plt.subplot(2,1,1)
        plt.imshow(gray2rgb(np.array(Image.open(file_name))))

        plt.subplot(2,1,2)
        image_3d = self.load_image(id)
        image_3d = rotate(image_3d, 1, resize=True)
        image_3d = self.fit_image(image_3d)
        plt.imshow(image_3d)

        
    def encode_all(self):
        for i in range(self.train_matrix.shape[0]):
            id = self.train_matrix[i, 0]
            label = self.train_matrix[i, 1]            

            self.encode_train(id, label)

        for i in range(self.test_matrix.shape[0]):
            id = self.test_matrix[i]
            self.encode_test(id)
            
    def encode_train(self, id, label):
        image = self.load_image(id)

        for angle in [0, 0.5, 1, -1, -0.5]:
        # for angle in (1, 2, 3, 4, 5, 90, 180, 270, -5, -4, -3, -2, -1):
            image = rotate(image, angle, resize=True)
            image = self.fit_image(image)
            
            i = np.random.randint(10)
            if angle == 0 or i > 0:
                path = os.path.join("./tfr_train/",  str(id) + "_" + str(angle) + ".trf")    
            else:
                path = os.path.join("./tfr_validation/", str(id) + "_" + str(angle) + ".trf")    
        
            with tf.python_io.TFRecordWriter(path) as writer:
                image_raw = image.tostring()

                example = tf.train.Example(features=tf.train.Features(feature={
                        'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
                        'image_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_raw]))}))

                writer.write(example.SerializeToString())

    def encode_test(self, id):
        image_3d = self.load_image(id)
        image_3d = self.fit_image(image_3d)
        path = os.path.join("./tfr_test/",  str(id) + ".trf")
        
        with tf.python_io.TFRecordWriter(path) as writer:
            image_raw = image_3d.tostring()
            example = tf.train.Example(features=tf.train.Features(feature={
                    'id': tf.train.Feature(int64_list=tf.train.Int64List(value=[id])),
                    'image_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_raw]))}))

            writer.write(example.SerializeToString())

gen = TFRecordGenerator(image_size=IMAGE_SIZE)
gen.encode_all()