（Tensorflow学习）药品的缺陷分类应用 tensorflow实现

在17版本的halcon中有深度学习的例子，关于分类的。

如下面三类，脏污，裂纹，和好的。

这边可以用tensorflow重现下，因为在笔记本上跑的，显卡不行，把原来300*300 resize到150*150

网络搭建个简单的，两层卷积就行，用keras实现。

读取图像用keras里的那个flow_from_directory，非常方便，最后保存h5模型，方便使用

训练源码

from keras.models import Sequential
from keras.layers import Conv2D, MaxPool2D, Flatten, Dropout, Dense
from keras.losses import categorical_crossentropy
from keras.optimizers import Adadelta
from keras.preprocessing.image import ImageDataGenerator

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.8
set_session(tf.Session(config=config))

model=Sequential()
model.add(Conv2D(32,(5,5),activation='relu',input_shape=[150,150,3]))
model.add(MaxPool2D(pool_size=(2,2),strides=[2,2],padding='SAME'))
model.add(Conv2D(64,(5,5),activation='relu'))
model.add(MaxPool2D(pool_size=(2,2),strides=[1,1],padding='SAME'))
model.add(Flatten())
model.add(Dropout(0.5))
model.add(Dense(128,activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(3,activation='softmax'))
model.compile(loss=categorical_crossentropy,optimizer=Adadelta(),metrics=['accuracy'])


datagen = ImageDataGenerator(rescale=1./255)
train_generator = datagen.flow_from_directory(
    './pill/magnesium',
    target_size=(150, 150),
    batch_size=16)
model.fit_generator(train_generator, steps_per_epoch=100, epochs=50)

mp = "model.h5"
model.save(mp)

测试源码

from keras.models import load_model
from skimage import io,transform
import numpy as np

load_model=load_model('model.h5')
img=io.imread("a.png")
img=transform.resize(img,(150,150))
img3=img[np.newaxis,:,:,:]
predicted = load_model.predict(img3)
print(predicted)
io.imshow(img)
io.show()

OK