classifier2.py

import tensorflow as tf

import os, sys
import cv2, numpy as np
from random import shuffle

VALID=0
TRAIN=1

class Classifier():

	# initialize Classifier parameters	
	def __init__(self, ncls, shape, batch=16):
		self.inp_shape = shape
		self.batch_size = batch
		self.n_classes = ncls

		self.images = tf.placeholder(tf.float32, [batch, shape[0], shape[1], shape[2]])
		self.labels = tf.placeholder(tf.float32, [batch, ncls])

		self.dataset = []

	# load data from directory, given path and category
	def data_from_dir(self, path, cat):
		filenames = os.listdir(path)
		shuffle(filenames)
		size = len(filenames)
		trn_size = size*0.8
		for f in range(size):
			for c in range(len(cat)):
				if cat[c] in filenames[f]:
					if f <= trn_size:
						self.dataset.append((os.path.join(path, filenames[f]), c, 1))
					else:
						self.dataset.append((os.path.join(path, filenames[f]), c, 0))


	# load next batch of data.
	def next_batch(self, pos, mode):
		batch = self.batch_size

		imgs = np.zeros([batch, self.inp_shape[0], self.inp_shape[1], self.inp_shape[2]])
		lbls = np.zeros([batch, self.n_classes])
		b = 0 # iterator for batch
		while batch > 0:
			# batch over-flow, start from 0
			if pos >= len(self.dataset):
				pos = 0

			data = self.dataset[pos]
			if data[-1] == mode:
				img = cv2.imread(data[0])
				img = cv2.resize(img, (self.inp_shape[0], self.inp_shape[1]))/255.
				imgs[b, ...] = img
				# one-hot encoding
				lbls[b][data[1]] = 1.0
				batch -= 1
				b += 1

			pos += 1

		# return batch images, labels and current position
		return imgs, lbls, pos


	def model(self):
		# 1st Block
		conv1 = tf.layers.conv2d(self.images, 64, [3, 3], activation='relu', name='conv1')
		conv2 = tf.layers.conv2d(conv1, 64, [3, 3], activation='relu', name='conv2')
		pool2 = tf.layers.max_pooling2d(conv2, [2, 2], 2, name='pool2')

		# 2nd Block
		conv3 = tf.layers.conv2d(pool2, 128, [3, 3], activation='relu', name='conv3')
		conv4 = tf.layers.conv2d(conv3, 128, [3, 3], activation='relu', name='conv4')
		pool4 = tf.layers.max_pooling2d(conv4, [2, 2], 2, name='pool4')

		# 3rd Block
		conv5 = tf.layers.conv2d(pool4, 256, [3, 3], activation='relu', name='conv5')
		conv6 = tf.layers.conv2d(conv5, 256, [3, 3], activation='relu', name='conv6')
		conv7 = tf.layers.conv2d(conv6, 256, [3, 3], activation='relu', name='conv7')
		pool7 = tf.layers.max_pooling2d(conv7, [2, 2], 2, name='pool7')

		# 4th Block
		conv8 = tf.layers.conv2d(pool7, 512, [3, 3], activation='relu', name='conv8')
		conv9 = tf.layers.conv2d(conv8, 512, [3, 3], activation='relu', name='conv9')
		conv10 = tf.layers.conv2d(conv9, 512, [3, 3], activation='relu', name='conv10')
		pool10 = tf.layers.max_pooling2d(conv10, [2, 2], 2, name='pool10')

		# 5th Block
		conv11 = tf.layers.conv2d(pool10, 512, [3, 3], activation='relu', name='conv11')
		conv12 = tf.layers.conv2d(conv11, 512, [3, 3], activation='relu', name='conv12')
		conv13 = tf.layers.conv2d(conv12, 512, [3, 3], activation='relu', name='conv13')
		pool13 = tf.layers.max_pooling2d(conv13, [2, 2], 2, name='pool13')

		fc = tf.layers.Flatten()(pool13)
		d1 = tf.layers.dense(fc, 4096)
		d2 = tf.layers.dense(d1, 4096)
		out = tf.layers.dense(d2, self.n_classes, name='out')

		return out


	def train(self, lr=0.01, itern=1000):
		self.lr_rate = lr
		self.n_itern = itern

		y = self.model()

		init = tf.global_variables_initializer()
		sess = tf.Session()
		sess.run(init)

		loss_fn = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=y, labels=self.labels))
		opt = tf.train.GradientDescentOptimizer(self.lr_rate).minimize(loss_fn)

		loss_trn, loss_val = 0, 0
		post, posv = 0, 0
		for i in range(self.n_itern):
			X, Y, post = self.next_batch(post, TRAIN)
			train, loss_trn = sess.run([opt, loss_fn], feed_dict={self.images: X, self.labels: Y})

			X, Y, posv = self.next_batch(posv, VALID)			
			loss_val = sess.run(loss_fn, feed_dict={self.images: X, self.labels: Y})

			print ('>> itern = %d, train_loss = %.3f, valid_loss = %.3f, lr=%.4f' % (i, loss_trn, loss_val, self.lr_rate))

		return loss_trn, loss_val


def main():
	inp_path = sys.argv[1]

	clf = Classifier(2, (224, 224, 3), 8)
	clf.data_from_dir(inp_path, ['cat', 'dog'])

	loss1, loss2 = clf.train(0.01)


if __name__=="__main__":
	main()