ypwhs / dogs_vs_cats Goto Github PK

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'
from keras.models import Model
from keras.callbacks import ReduceLROnPlateau
from keras.utils import plot_model
from keras.layers import *
from keras.applications import *
from keras import backend as K
from keras.preprocessing.image import *
import matplotlib.pyplot as plt
import h5py

base_models = locals()
use_models = {#ResNet50:{'preprocessing':None, 'size':(224, 224)},
              InceptionV3:{'preprocessing':inception_v3.preprocess_input, 'size':(299, 299), 'device':'/gpu:0'},
              Xception:{'preprocessing':xception.preprocess_input, 'size':(299, 299), 'device':'/gpu:1'}, 
              InceptionResNetV2:{'preprocessing':inception_resnet_v2.preprocess_input, 'size':(299, 299), 'device':'/gpu:2'}
              }

model_bottlenecks = []
input_tensor = Input((299, 299, 3))
for model, method in use_models.items():
    width, height = method['size']
    preprocessing = method['preprocessing']
    device = method['device']
    # input_tensor = Input((height, width, 3))
    if preprocessing:
        x = Lambda(preprocessing)(input_tensor)
    with K.tf.device(device):
      base_models[model.__name__] = model(input_tensor=x, weights='imagenet', include_top=False)
      for layer in base_models[model.__name__] .layers:
          layer.trainable = False
      model_bottleneck = GlobalAveragePooling2D()(base_models[model.__name__].output)
      model_bottlenecks.append(model_bottleneck)

with K.tf.device('/gpu:0'):
  stack_bottleneck = concatenate(model_bottlenecks)
  fc = Dropout(0.5)(stack_bottleneck)
  fc = Dense(120, activation='softmax')(fc)

stack_model = Model(input_tensor, fc)
# plot_model(stack_model, to_file='model.png')


datagen = ImageDataGenerator(rotation_range=10,
                             width_shift_range=0.1,
                             height_shift_range=0.1,
                             zoom_range=0.1,
                             horizontal_flip=True,
                             vertical_flip=True
                             )
train_generator = datagen.flow_from_directory("raw_images", (299, 299), batch_size=64)
test_generator = datagen.flow_from_directory("test", (299, 299), batch_size=1)

stack_model.compile(optimizer='rmsprop',
              loss='categorical_crossentropy',
              metrics=['accuracy'])
learning_rate_reduction = ReduceLROnPlateau(monitor='loss', 
                                            patience=3, 
                                            verbose=1, 
                                            factor=0.5, 
                                            min_lr=1e-8)

stack_model.fit_generator(train_generator,
                    steps_per_epoch=train_generator.n/64, 
                    epochs=100, 
                    verbose=1,
                    callbacks=[learning_rate_reduction])
stack_model.save_weights('stack_model.h5')

# coding: utf-8

# In[1]:

import os
import shutil

# In[2]:
SMALL_DATASET_COUNT = 1000

train_filenames = os.listdir('train')
train_cat = filter(lambda x: x[:3] == 'cat', train_filenames)
train_dog = filter(lambda x: x[:3] == 'dog', train_filenames)

test_file_dir = os.listdir('test')
test_filenames = filter(lambda x: int(x[:-4]) <= SMALL_DATASET_COUNT, test_file_dir)


# In[3]:

def rmrf_mkdir(dirname):
    if os.path.exists(dirname):
        shutil.rmtree(dirname)
    os.mkdir(dirname)


# rmrf_mkdir('train2')
# os.mkdir('train2/cat')
# os.mkdir('train2/dog')
#
# rmrf_mkdir('test2')
# os.symlink('../test/', 'test2/test')
#
#
# for filename in train_cat:
#     os.symlink('../../train/'+filename, 'train2/cat/'+filename)
#
# for filename in train_dog:
#     os.symlink('../../train/'+filename, 'train2/dog/'+filename)


# In[ ]:


rmrf_mkdir('train-small-dataset')
os.mkdir('train-small-dataset/cat')
os.mkdir('train-small-dataset/dog')

file_count = 0
for filename in train_cat:
    os.symlink('../../train/' + filename, 'train-small-dataset/cat/' + filename)
    file_count += 1
    if file_count >= SMALL_DATASET_COUNT:
        break

file_count = 0
for filename in train_dog:
    os.symlink('../../train/' + filename, 'train-small-dataset/dog/' + filename)
    file_count += 1
    if file_count >= SMALL_DATASET_COUNT:
        break

rmrf_mkdir('test-small-dataset')
rmrf_mkdir('test-small-dataset/test')

file_count = 0
for filename in test_filenames:
    os.symlink('../../test/' + filename, 'test-small-dataset/test/' + filename)
    file_count += 1
    if file_count >= SMALL_DATASET_COUNT:
        break