Wrong predicted images size from predict_generator for multi label Unet

问题

Im using aerial images to segment road and centerline using multi label u-net, my test generator is look like this

def testGenerator(test_path= "data\\membrane\\test\\image",num_image = 1584,target_size = (224,224),flag_multi_class = False,as_gray = False):
for i in range(num_image):
    img = io.imread(os.path.join(test_path,"%d.jpg"%i),as_gray = as_gray)
    img = img / 255.
    img = trans.resize(img,target_size)
    img = np.reshape(img,img.shape) if (not flag_multi_class) else img
    img = np.reshape(img,(1,)+img.shape)
    yield img

based on Unet

when I'm trying to get the prediction from predict_generator as shown in the code below

 testGene = testGenerator("data\\membrane\\test\\image")results =model.predict_generator(testGene,1584,verbose=1)saveResult("data\\membrane\\tes\\results",results)`

I used for visualization this part of code:

def labelVisualize(num_class,color_dict,img):
    img = img[:,:,0] if len(img.shape) == 3 else img
    img_out = np.zeros(img.shape + (3,))
    for i in range(num_class):
        img_out[img == i,:] = color_dict[i]
    return img_out / 255

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.tif"%(i)),skimage.img_as_ubyte(img))

I faced three problems :

1- despite I have 1584 images as my test sample where batch_size= 1 Ive got only 2 predicted images one for each class( road and center line ) predict_generator process

2- my target size is(224,224) but for the prediction Ive got image size of (224,1584).

3- the prediction is unclear for both classes as shown in the images first image in test set, predicted_1, predicted_2

anyone has an idea where is my mistake please, any help will be appreciated, thank you in advance.

来源：https://stackoverflow.com/questions/60861264/wrong-predicted-images-size-from-predict-generator-for-multi-label-unet