Tensorflow U-Net 分割掩码输入

Question

I am new to tensorflow and Semantic segmentation.

I am designing a U-Net for semantic segmentaion. Each image has one object that I want to classify. But in total I have images of 10 different objects. I am confused, how can I prepare my mask input? Is it considered as multi-label segmentation or only for one class?

Should I convert my input to one hot encoded? Should I use to_categorical? I find exaples for multi-class segmentation, but I don't know, If that's the case here. Because in one image I only have one object to detect/classify.

I tried using this as my code for input. But I am not sure, what I am doing is right or not.

#Generation of batches of image and mask
class DataGen(keras.utils.Sequence):
    def __init__(self, image_names, path, batch_size, image_size=128):
        self.image_names = image_names
        self.path = path
        self.batch_size = batch_size
        self.image_size = image_size

    def __load__(self, image_name):
        # Path
        image_path = os.path.join(self.path, "images/aug_test", image_name) + ".png"
        mask_path = os.path.join(self.path, "masks/aug_test",image_name) +  ".png"

        # Reading Image
        image = cv2.imread(image_path, 1)
        image = cv2.resize(image, (self.image_size, self.image_size))


        # Reading Mask
        mask = cv2.imread(mask_path, -1)
        mask = cv2.resize(mask, (self.image_size, self.image_size))

        ## Normalizaing 
        image = image/255.0
        mask = mask/255.0

        return image, mask

    def __getitem__(self, index):
        if(index+1)*self.batch_size > len(self.image_names):
            self.batch_size = len(self.image_names) - index*self.batch_size

        image_batch = self.image_names[index*self.batch_size : (index+1)*self.batch_size]

        image = []
        mask  = []

        for image_name in image_batch:
            _img, _mask = self.__load__(image_name)
            image.append(_img)
            mask.append(_mask)

        #This is where I am defining my input
        image = np.array(image)
        mask  = np.array(mask)
        mask = tf.keras.utils.to_categorical(mask, num_classes=10, dtype='float32') #Is this true?


        return image, mask

    def __len__(self):
        return int(np.ceil(len(self.image_names)/float(self.batch_size)))

Is this true? If it is, then, to get the label/class as output what should I change in my input? Should I change the value of pixel of my mask according to my class?

Here is my U-Net architecture.

# Convolution and deconvolution Blocks

def down_scaling_block(x, filters, kernel_size=(3, 3), padding="same", strides=1):
    conv = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation="relu")(x)
    conv = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation="relu")(conv)
    pool = keras.layers.MaxPool2D((2, 2), (2, 2))(conv)
    return conv, pool

def up_scaling_block(x, skip, filters, kernel_size=(3, 3), padding="same", strides=1):
    conv_t = keras.layers.UpSampling2D((2, 2))(x)
    concat = keras.layers.Concatenate()([conv_t, skip])
    conv = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation="relu")(concat)
    conv = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation="relu")(conv)
    return conv

def bottleneck(x, filters, kernel_size=(3, 3), padding="same", strides=1):
    conv = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation="relu")(x)
    conv = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation="relu")(conv)
    return conv

def UNet():
    filters = [16, 32, 64, 128, 256]
    inputs = keras.layers.Input((image_size, image_size, 3))

    '''inputs2 = keras.layers.Input((image_size, image_size, 1))
       conv1_2, pool1_2 = down_scaling_block(inputs2, filters[0])'''

    Input = inputs
    conv1, pool1 = down_scaling_block(Input, filters[0])
    conv2, pool2 = down_scaling_block(pool1, filters[1])
    conv3, pool3 = down_scaling_block(pool2, filters[2])
    '''conv3 = keras.layers.Conv2D(filters[2], kernel_size=(3,3), padding="same", strides=1, activation="relu")(pool2)
    conv3 = keras.layers.Conv2D(filters[2], kernel_size=(3,3), padding="same", strides=1, activation="relu")(conv3)
    drop3 = keras.layers.Dropout(0.5)(conv3)
    pool3 = keras.layers.MaxPooling2D((2,2), (2,2))(drop3)'''
    conv4, pool4 = down_scaling_block(pool3, filters[3])

    bn = bottleneck(pool4, filters[4])

    deConv1 = up_scaling_block(bn, conv4, filters[3]) #8 -> 16
    deConv2 = up_scaling_block(deConv1, conv3, filters[2]) #16 -> 32
    deConv3 = up_scaling_block(deConv2, conv2, filters[1]) #32 -> 64
    deConv4 = up_scaling_block(deConv3, conv1, filters[0]) #64 -> 128

    outputs = keras.layers.Conv2D(10, (1, 1), padding="same", activation="softmax")(deConv4)
    model = keras.models.Model(inputs, outputs)
    return model

model = UNet()
model.compile(optimizer='adam', loss="categorical_crossentropy", metrics=["acc"])

train_gen = DataGen(train_img, train_path, image_size=image_size, batch_size=batch_size)
valid_gen = DataGen(valid_img, train_path, image_size=image_size, batch_size=batch_size)
test_gen = DataGen(test_img, test_path, image_size=image_size, batch_size=batch_size)

train_steps = len(train_img)//batch_size
valid_steps = len(valid_img)//batch_size

model.fit_generator(train_gen, validation_data=valid_gen, steps_per_epoch=train_steps, validation_steps=valid_steps, 
                    epochs=epochs)

I hope that I explained my question properly. Any help appriciated!

UPDATE: I changed the value of each pixel in mask as per object class. (If the image contains object which I want to classify as object no. 2, then I changed the value of mask pixel to 2. the whole array of mask will contain 0(bg) and 2(object). Accordingly for each object, the mask will contain 0 and 3, 0 and 10 etc.)

Here I first changed the mask to binary and then if the value of pixel is greater than 1, I changed it to 1 or 2 or 3. (according to object/class no.)

Then I converted them to one_hot with to_categorical as shown in my code. training runs but the network doesnt learn anything. Accuracy and loss keep swinging between two values. What is my mistake here? Am I making a mistake at generating mask (changing the value of pixels?) Or at the function to_categorical?

PROBLEM FOUND: I was making an error while creating mask.. I was reading image with cv2, which reads image as heightxwidth.. I was creating mask with pixel values according to class, after considering my image dimention as widthxheight.. Which was causing problem and making network not to learn anything.. It is working now..

Answer 1

Each image has one object that I want to classify. But in total I have images of 10 different objects. I am confused, how can I prepare my mask input? Is it considered as multi-label segmentation or only for one class?

If your dataset has N different labels (ie: 0 - background, 1 - dogs, 2 -cats...), you have a multi class problem, even if your images contain only kind of object.

Should I convert my input to one hot encoded? Should I use to_categorical?

Yes, you should one-hot encode your labels. Using to_categorical boils down to the source format of your labels. Say you have N classes and your labels are (height, width, 1), where each pixel has a value in range [0,N). In that case keras.utils.to_categorical(label, N) will provide a float (height,width,N) label, where each pixel is 0 or 1. And you don't have to divide by 255.

if your source format is different, you may have to use a custom function to get the same output format.

Check out this repo (not my work): keras-unet . The notebooks folder contain two examples to train a u-net on small datasets. They are not multiclass, but it is easy to go step by step to use your own dataset. Star by loading your labels as:

im = Image.open(mask).resize((512,512))
im = to_categorical(im,NCLASSES)

reshape and normalize like this:

x = np.asarray(imgs_np, dtype=np.float32)/255
y = np.asarray(masks_np, dtype=np.float32)
y = y.reshape(y.shape[0], y.shape[1], y.shape[2], NCLASSES)
x = x.reshape(x.shape[0], x.shape[1], x.shape[2], 3)

adapt your model to NCLASSES

model = custom_unet(
input_shape,
use_batch_norm=False,
num_classes=NCLASSES,
filters=64,
dropout=0.2,
output_activation='softmax')

select the correct loss:

from keras.losses import categorical_crossentropy
model.compile(    
   optimizer=SGD(lr=0.01, momentum=0.99),
   loss='categorical_crossentropy',    
   metrics=[iou, iou_thresholded])

Hope it helps