Created
March 6, 2019 06:44
-
-
Save danishansari/8406afd287e64759fd3000e841147351 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment