Tiny network modifications

README.md

@@ -4,6 +4,11 @@ The architecture was inspired by [U-Net: Convolutional Networks for Biomedical I
 
 ---
 
+Lab work
+
+Modified from zhixuhao version, the current version can be directly deployed locally, debugs the network, etc....
+
+
 ## Overview
 
 ### Data

data.py

@@ -5,6 +5,7 @@
 import glob
 import skimage.io as io
 import skimage.transform as trans
+import cv2
 
 Sky = [128,128,128]
 Building = [128,0,0]
@@ -119,6 +120,13 @@ def labelVisualize(num_class,color_dict,img):
 
 
 def saveResult(save_path,npyfile,flag_multi_class = False,num_class = 2):
-    for i,item in enumerate(npyfile):
-        img = labelVisualize(num_class,COLOR_DICT,item) if flag_multi_class else item[:,:,0]
-        io.imsave(os.path.join(save_path,"%d_predict.png"%i),img)
+    for i, item in enumerate(npyfile):
+        img = labelVisualize(num_class, COLOR_DICT, item) if flag_multi_class else item[:, :, 0]
+        #print(img)  ## ex. [[0.16509348 0.08227982 ... 0.07113015]
+        img = (img > 0.5).astype(np.uint8)  # .reshape(256, 256)
+        #print(img)  ## ex. [[ 0 0 0 0 ... 1 1 1 ]]
+        img = (img * 255)
+        #print(img)  ## ex. [[0 0 00 ... 255 255 255]]
+        # img = labelVisualize(num_class,COLOR_DICT,item) if flag_multi_class else item[:,:,0]
+        # io.imsave(os.path.join(save_path,"%d_predict.png"%i),img)
+        cv2.imwrite(os.path.join(save_path, "%d_predict.png" % i), img)

main.py

@@ -1,5 +1,7 @@
 from model import *
 from data import *
+import warnings
+warnings.filterwarnings('ignore')
 
 #os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
@@ -15,8 +17,7 @@
 
 model = unet()
 model_checkpoint = ModelCheckpoint('unet_membrane.hdf5', monitor='loss',verbose=1, save_best_only=True)
-model.fit_generator(myGene,steps_per_epoch=300,epochs=1,callbacks=[model_checkpoint])
-
+model.fit_generator(myGene,steps_per_epoch=300,epochs=5,callbacks=[model_checkpoint])
 testGene = testGenerator("data/membrane/test")
 results = model.predict_generator(testGene,30,verbose=1)
-saveResult("data/membrane/test",results)
+saveResult("data/membrane/result",results)

model.py

@@ -52,7 +52,7 @@ def unet(pretrained_weights = None,input_size = (256,256,1)):
     conv9 = Conv2D(2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
     conv10 = Conv2D(1, 1, activation = 'sigmoid')(conv9)
 
-    model = Model(input = inputs, output = conv10)
+    model = Model(inputs,conv10)
 
     model.compile(optimizer = Adam(lr = 1e-4), loss = 'binary_crossentropy', metrics = ['accuracy'])