model.py 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. import os
  2. from keras.preprocessing import image
  3. import matplotlib.pyplot as plt
  4. import numpy as np
  5. from keras.utils.np_utils import to_categorical
  6. import random,shutil
  7. from keras.models import Sequential
  8. from keras.layers import Dropout,Conv2D,Flatten,Dense, MaxPooling2D, BatchNormalization
  9. from keras.models import load_model
  10. def generator(dir, gen=image.ImageDataGenerator(rescale=1./255), shuffle=True,batch_size=1,target_size=(24,24),class_mode='categorical' ):
  11. return gen.flow_from_directory(dir,batch_size=batch_size,shuffle=shuffle,color_mode='grayscale',class_mode=class_mode,target_size=target_size)
  12. BS= 32
  13. TS=(24,24)
  14. train_batch= generator('data/train',shuffle=True, batch_size=BS,target_size=TS)
  15. valid_batch= generator('data/valid',shuffle=True, batch_size=BS,target_size=TS)
  16. SPE= len(train_batch.classes)//BS
  17. VS = len(valid_batch.classes)//BS
  18. print(SPE,VS)
  19. # img,labels= next(train_batch)
  20. # print(img.shape)
  21. model = Sequential([
  22. Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=(24,24,1)),
  23. MaxPooling2D(pool_size=(1,1)),
  24. Conv2D(32,(3,3),activation='relu'),
  25. MaxPooling2D(pool_size=(1,1)),
  26. #32 convolution filters used each of size 3x3
  27. #again
  28. Conv2D(64, (3, 3), activation='relu'),
  29. MaxPooling2D(pool_size=(1,1)),
  30. #64 convolution filters used each of size 3x3
  31. #choose the best features via pooling
  32. #randomly turn neurons on and off to improve convergence
  33. Dropout(0.25),
  34. #flatten since too many dimensions, we only want a classification output
  35. Flatten(),
  36. #fully connected to get all relevant data
  37. Dense(128, activation='relu'),
  38. #one more dropout for convergence' sake :)
  39. Dropout(0.5),
  40. #output a softmax to squash the matrix into output probabilities
  41. Dense(2, activation='softmax')
  42. ])
  43. model.compile(optimizer='adam',loss='categorical_crossentropy',metrics=['accuracy'])
  44. model.fit_generator(train_batch, validation_data=valid_batch,epochs=15,steps_per_epoch=SPE ,validation_steps=VS)
  45. model.save('models/cnnCat2.h5', overwrite=True)