{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#data.py  \nfrom __future__ import print_function\nfrom keras.preprocessing.image import ImageDataGenerator\nimport numpy as np \nimport os\nimport glob\nimport skimage.io as io\nimport skimage.transform as trans\n\nSky = [128,128,128]\nBuilding = [128,0,0]\nPole = [192,192,128]\nRoad = [128,64,128]\nPavement = [60,40,222]\nTree = [128,128,0]\nSignSymbol = [192,128,128]\nFence = [64,64,128]\nCar = [64,0,128]\nPedestrian = [64,64,0]\nBicyclist = [0,128,192]\nUnlabelled = [0,0,0]\n\nCOLOR_DICT = np.array([Sky, Building, Pole, Road, Pavement,\n                          Tree, SignSymbol, Fence, Car, Pedestrian, Bicyclist, Unlabelled])\n\n\ndef adjustData(img,mask,flag_multi_class,num_class):\n    if(flag_multi_class):\n        img = img / 255\n        mask = mask[:,:,:,0] if(len(mask.shape) == 4) else mask[:,:,0]\n        new_mask = np.zeros(mask.shape + (num_class,))\n        for i in range(num_class):\n            #for one pixel in the image, find the class in mask and convert it into one-hot vector\n            #index = np.where(mask == i)\n            #index_mask = (index[0],index[1],index[2],np.zeros(len(index[0]),dtype = np.int64) + i) if (len(mask.shape) == 4) else (index[0],index[1],np.zeros(len(index[0]),dtype = np.int64) + i)\n            #new_mask[index_mask] = 1\n            new_mask[mask == i,i] = 1\n        new_mask = np.reshape(new_mask,(new_mask.shape[0],new_mask.shape[1]*new_mask.shape[2],new_mask.shape[3])) if flag_multi_class else np.reshape(new_mask,(new_mask.shape[0]*new_mask.shape[1],new_mask.shape[2]))\n        mask = new_mask\n    elif(np.max(img) > 1):\n        img = img / 255\n        mask = mask /255\n        mask[mask > 0.5] = 1\n        mask[mask <= 0.5] = 0\n    return (img,mask)\n\n\n\ndef trainGenerator(batch_size,train_path,image_folder,mask_folder,aug_dict,image_color_mode = \"grayscale\",\n                    mask_color_mode = \"grayscale\",image_save_prefix  = \"image\",mask_save_prefix  = \"mask\",\n                    flag_multi_class = False,num_class = 2,save_to_dir = None,target_size = (256,256),seed = 1):\n    '''\n    can generate image and mask at the same time\n    use the same seed for image_datagen and mask_datagen to ensure the transformation for image and mask is the same\n    if you want to visualize the results of generator, set save_to_dir = \"your path\"\n    '''\n    image_datagen = ImageDataGenerator(**aug_dict)\n    mask_datagen = ImageDataGenerator(**aug_dict)\n    image_generator = image_datagen.flow_from_directory(\n        train_path,\n        classes = [image_folder],\n        class_mode = None,\n        color_mode = image_color_mode,\n        target_size = target_size,\n        batch_size = batch_size,\n        save_to_dir = save_to_dir,\n        save_prefix  = image_save_prefix,\n        seed = seed)\n    mask_generator = mask_datagen.flow_from_directory(\n        train_path,\n        classes = [mask_folder],\n        class_mode = None,\n        color_mode = mask_color_mode,\n        target_size = target_size,\n        batch_size = batch_size,\n        save_to_dir = save_to_dir,\n        save_prefix  = mask_save_prefix,\n        seed = seed)\n    train_generator = zip(image_generator, mask_generator)\n    for (img,mask) in train_generator:\n        img,mask = adjustData(img,mask,flag_multi_class,num_class)\n        yield (img,mask)\n\n\n\ndef testGenerator(test_path,num_image = 30,target_size = (256,256),flag_multi_class = False,as_gray = True):\n    for i in range(num_image):\n        img = io.imread(os.path.join(test_path,\"%d.png\"%i),as_gray = as_gray)\n        img = img / 255\n        img = trans.resize(img,target_size)\n        img = np.reshape(img,img.shape+(1,)) if (not flag_multi_class) else img\n        img = np.reshape(img,(1,)+img.shape)\n        yield img\n\n\ndef geneTrainNpy(image_path,mask_path,flag_multi_class = False,num_class = 2,image_prefix = \"image\",mask_prefix = \"mask\",image_as_gray = True,mask_as_gray = True):\n    image_name_arr = glob.glob(os.path.join(image_path,\"%s*.png\"%image_prefix))\n    image_arr = []\n    mask_arr = []\n    for index,item in enumerate(image_name_arr):\n        img = io.imread(item,as_gray = image_as_gray)\n        img = np.reshape(img,img.shape + (1,)) if image_as_gray else img\n        mask = io.imread(item.replace(image_path,mask_path).replace(image_prefix,mask_prefix),as_gray = mask_as_gray)\n        mask = np.reshape(mask,mask.shape + (1,)) if mask_as_gray else mask\n        img,mask = adjustData(img,mask,flag_multi_class,num_class)\n        image_arr.append(img)\n        mask_arr.append(mask)\n    image_arr = np.array(image_arr)\n    mask_arr = np.array(mask_arr)\n    return image_arr,mask_arr\n\n\ndef labelVisualize(num_class,color_dict,img):\n    img = img[:,:,0] if len(img.shape) == 3 else img\n    img_out = np.zeros(img.shape + (3,))\n    for i in range(num_class):\n        img_out[img == i,:] = color_dict[i]\n    return img_out / 255\n\n\n\ndef saveResult(save_path,npyfile,flag_multi_class = False,num_class = 2):\n    for i,item in enumerate(npyfile):\n        img = labelVisualize(num_class,COLOR_DICT,item) if flag_multi_class else item[:,:,0]\n        io.imsave(os.path.join(save_path,\"%d_predict.png\"%i),img)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#dataprepare.py\nfrom data import *\n\n#if you don't want to do data augmentation, set data_gen_args as an empty dict.\n#data_gen_args = dict()\n\ndata_gen_args = dict(rotation_range=0.2,\n                    width_shift_range=0.05,\n                    height_shift_range=0.05,\n                    shear_range=0.05,\n                    zoom_range=0.05,\n                    horizontal_flip=True,\n                    fill_mode='nearest')\nmyGenerator = trainGenerator(20,'data/membrane/train','image','label',data_gen_args,save_to_dir = \"data/membrane/train/aug\")\n\n\n#you will see 60 transformed images and their masks in data/membrane/train/aug\nnum_batch = 3\nfor i,batch in enumerate(myGenerator):\n    if(i >= num_batch):\n        break\n\nimage_arr,mask_arr = geneTrainNpy(\"data/membrane/train/aug/\",\"data/membrane/train/aug/\")\n#np.save(\"data/image_arr.npy\",image_arr)\n#np.save(\"data/mask_arr.npy\",mask_arr)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#main.py\nimport numpy as np\nimport cv2\nfrom matplotlib import pyplot as plt\n\ndef sobelOperator(img):\n    container = np.copy(img)\n    size = container.shape\n    for i in range(1, size[0] - 1):\n        for j in range(1, size[1] - 1):\n            gx = (img[i - 1][j - 1] + 2*img[i][j - 1] + img[i + 1][j - 1]) - (img[i - 1][j + 1] + 2*img[i][j + 1] + img[i + 1][j + 1])\n            gy = (img[i - 1][j - 1] + 2*img[i - 1][j] + img[i - 1][j + 1]) - (img[i + 1][j - 1] + 2*img[i + 1][j] + img[i + 1][j + 1])\n            container[i][j] = min(255, np.sqrt(gx**2 + gy**2))\n    return container\n    pass\n\n\n#test test_img\n\nimg = cv2.cvtColor(cv2.imread(\"test_img.png\"), cv2.COLOR_BGR2GRAY)\nimg = sobelOperator(img)\nimg = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)\nplt.imshow(img)\nplt.show()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#model.py\nimport numpy as np \nimport os\nimport skimage.io as io\nimport skimage.transform as trans\nimport numpy as np\nfrom keras.models import *\nfrom keras.layers import *\nfrom keras.optimizers import *\nfrom keras.callbacks import ModelCheckpoint, LearningRateScheduler\nfrom keras import backend as keras\n\n\ndef unet(pretrained_weights = None,input_size = (256,256,1)):\n    inputs = Input(input_size)\n    conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(inputs)\n    conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv1)\n    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)\n    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool1)\n    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv2)\n    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)\n    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool2)\n    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv3)\n    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)\n    conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool3)\n    conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv4)\n    drop4 = Dropout(0.5)(conv4)\n    pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)\n\n    conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool4)\n    conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv5)\n    drop5 = Dropout(0.5)(conv5)\n\n    up6 = Conv2D(512, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(drop5))\n    merge6 = concatenate([drop4,up6], axis = 3)\n    conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6)\n    conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6)\n\n    up7 = Conv2D(256, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv6))\n    merge7 = concatenate([conv3,up7], axis = 3)\n    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge7)\n    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)\n\n    up8 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv7))\n    merge8 = concatenate([conv2,up8], axis = 3)\n    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge8)\n    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv8)\n\n    up9 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv8))\n    merge9 = concatenate([conv1,up9], axis = 3)\n    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge9)\n    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)\n    conv9 = Conv2D(2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)\n    conv10 = Conv2D(1, 1, activation = 'sigmoid')(conv9)\n\n    model = Model(input = inputs, output = conv10)\n\n    model.compile(optimizer = Adam(lr = 1e-4), loss = 'binary_crossentropy', metrics = ['accuracy'])\n    \n    #model.summary()\n\n    if(pretrained_weights):\n    \tmodel.load_weights(pretrained_weights)\n\n    return model\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nfrom model import *\nfrom data import *\n\ndata_gen_args = dict(rotation_range=0.2,\n                    width_shift_range=0.05,\n                    height_shift_range=0.05,\n                    shear_range=0.05,\n                    zoom_range=0.05,\n                    horizontal_flip=True,\n                    fill_mode='nearest')\nmyGene = trainGenerator(2,'data/membrane/train','image','label',data_gen_args,save_to_dir = None)\nmodel = unet()\nmodel_checkpoint = ModelCheckpoint('unet_membrane.hdf5', monitor='loss',verbose=1, save_best_only=True)\nmodel.fit_generator(myGene,steps_per_epoch=2000,epochs=5,callbacks=[model_checkpoint])\n\n\n#imgs_train,imgs_mask_train = geneTrainNpy(\"data/membrane/train/aug/\",\"data/membrane/train/aug/\")\n#model.fit(imgs_train, imgs_mask_train, batch_size=2, nb_epoch=10, verbose=1,validation_split=0.2, shuffle=True, callbacks=[model_checkpoint])\n\n\n\ntestGene = testGenerator(\"data/membrane/test\")\nmodel = unet()\nmodel.load_weights(\"unet_membrane.hdf5\")\nresults = model.predict_generator(testGene,30,verbose=1)\nsaveResult(\"data/membrane/test\",results)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport cv2\nfrom matplotlib import pyplot as plt\n\ndef sobelOperator(img):\n    container = np.copy(img)\n    size = container.shape\n    for i in range(1, size[0] - 1):\n        for j in range(1, size[1] - 1):\n            gx = (img[i - 1][j - 1] + 2*img[i][j - 1] + img[i + 1][j - 1]) - (img[i - 1][j + 1] + 2*img[i][j + 1] + img[i + 1][j + 1])\n            gy = (img[i - 1][j - 1] + 2*img[i - 1][j] + img[i - 1][j + 1]) - (img[i + 1][j - 1] + 2*img[i + 1][j] + img[i + 1][j + 1])\n            container[i][j] = min(255, np.sqrt(gx**2 + gy**2))\n    return container\n    pass\n\nimg = cv2.cvtColor(cv2.imread(\"test_img.png\"), cv2.COLOR_BGR2GRAY)\nimg = sobelOperator(img)\nimg = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)\nplt.imshow(img)\nplt.show()","metadata":{},"execution_count":null,"outputs":[]}]}