【Keras】基于SegNet和U-Net的遥感图像语义分割(三)

预测整张遥感影像：unet_predict.py

代码：
https://github.com/fuyou123/Segmentation_Unet

前期的工作是裁剪数据，得到训练集。并将训练集送入U_net模型进行训练，得到训练模型。

这里需要思考一下怎么预测整张遥感图像。我们知道，我们训练模型时选择的图片输入是256×256，所以我们预测时也要采用256×256的图片尺寸送进模型预测。现在我们要考虑一个问题，我们该怎么将这些预测好的小图重新拼接成一个大图呢？这里给出一个最基础的方案：先给大图做padding 0操作，得到一副padding过的大图，同时我们也生成一个与该图一样大的全0图A，把图像的尺寸补齐为256的倍数，然后以256为步长切割大图，依次将小图送进模型预测，预测好的小图则放在A的相应位置上，依次进行，最终得到预测好的整张大图（即A），再做图像切割，切割成原先图片的尺寸，完成整个预测流程。

1. args_parse（）# 参数设定

def args_parse():
# construct the argument parse and parse the arguments
    ap = argparse.ArgumentParser()
    ap.add_argument("-m", "--model", default="Trained_Unet_Model.h5",
        help="path to trained model ")
    ap.add_argument("-s", "--stride", type=int, default=image_size,
        help="crop slide stride")  # 滑动步长
    args = vars(ap.parse_args())    
    return args

2. predict(args)

def predict(args):
    # load the trained convolutional neural network
    print("[INFO] loading network...")
    model = load_model(args["model"])

    stride = args['stride']
    for n in range(len(TEST_SET)):
        path = TEST_SET[n]
        #load the image
        image = cv2.imread('./test/' + path)

        # calculate the value of padding
        h,w,_ = image.shape
        padding_h = (h//stride + 1) * stride    # " // "表示整数除法
        padding_w = (w//stride + 1) * stride
        padding_img = np.zeros((padding_h,padding_w,3),dtype=np.uint8)
        padding_img[0:h,0:w,:] = image[:,:,:]
        padding_img = padding_img.astype("float") / 255.0
        padding_img = img_to_array(padding_img)   # 从(h,w,c)到(c,h,w)
        print('src{}:{}'.format(path,padding_img.shape))

        # take out the small image to predict
        mask_whole = np.zeros((padding_h,padding_w),dtype=np.uint8)
        for i in range(padding_h//stride):
            for j in range(padding_w//stride):
                # 小图左上角坐标(Li,Lj )
                Li,Lj = (i*stride,j*stride)
                crop = padding_img[:3,Li:Li+image_size,Lj:Lj+image_size]   # small image
                _,ch,cw = crop.shape
                if ch != 256 or cw != 256:
                    print('invalid size!')
                    continue
                # predeict the small image
                crop = np.expand_dims(crop, axis=0) # (c,h,w)--->(1,c,h,w)
                pred = model.predict(crop,verbose=2)
                # print (np.unique(pred))
                pred = pred.reshape((256,256)).astype(np.uint8)
                #print 'pred:',pred.shape
                mask_whole[Li:Li+image_size,Lj:Lj+image_size] = pred[:,:]

        # save the final result
        cv2.imwrite(predict_path+"/"+path[:-4]+'.png',mask_whole[0:h,0:w])