Update unet.py

Author: zhixuhao

unet.py

@@ -1,10 +1,12 @@
+import os 
+#os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 import numpy as np
 from keras.models import *
 from keras.layers import Input, merge, Conv2D, MaxPooling2D, UpSampling2D, Dropout, Cropping2D
 from keras.optimizers import *
 from keras.callbacks import ModelCheckpoint, LearningRateScheduler
 from keras import backend as keras
-from data import dataProcess
+from data import *
 
 class myUnet(object):
 
@@ -155,16 +157,28 @@ def train(self):
 
 		model_checkpoint = ModelCheckpoint('unet.hdf5', monitor='loss',verbose=1, save_best_only=True)
 		print('Fitting model...')
-		model.fit(imgs_train, imgs_mask_train, batch_size=1, nb_epoch=10, verbose=1, shuffle=True, callbacks=[model_checkpoint])
+		model.fit(imgs_train, imgs_mask_train, batch_size=4, nb_epoch=10, verbose=1,validation_split=0.2, shuffle=True, callbacks=[model_checkpoint])
 
 		print('predict test data')
 		imgs_mask_test = model.predict(imgs_test, batch_size=1, verbose=1)
-		np.save('imgs_mask_test.npy', imgs_mask_test)
+		np.save('../results/imgs_mask_test.npy', imgs_mask_test)
+
+	def save_img(self):
+
+		print("array to image")
+		imgs = np.load('imgs_mask_test.npy')
+		for i in range(imgs.shape[0]):
+			img = imgs[i]
+			img = array_to_img(img)
+			img.save("../results/%d.jpg"%(i))
+
+
 
 
 if __name__ == '__main__':
 	myunet = myUnet()
 	myunet.train()
+	myunet.save_img()