qubvel / segmentation_models

Segmentation models with pretrained backbones. Keras and TensorFlow Keras.
MIT License
4.7k stars 1.03k forks source link

Custom masks processing problem #355

Open EtagiBI opened 4 years ago

EtagiBI commented 4 years ago

Hello,

According to this example, here's a way to detect one desired class in a binary segmentation problem:

...
CLASSES = ['car']
...
n_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1)  # case for binary and multiclass segmentation
...
train_dataset = Dataset(
    x_train_dir, 
    y_train_dir, 
    classes=CLASSES, 
    augmentation=get_training_augmentation(),
    preprocessing=get_preprocessing(preprocess_input),
)

Now let's take a look at Dataset object:

self.class_values = [self.CLASSES.index(cls.lower()) for cls in classes] <-- in our case we have [0] here

    def __getitem__(self, i):

        # read data
        image = cv2.imread(self.images_fps[i])
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(self.masks_fps[i], 0)

        # extract certain classes from mask (e.g. cars)
        masks = [(mask == v) for v in self.class_values] <-- ATTENTION!!!

Since class_values in this case is equal to [0], all elements of the mask matrix are compared to 0 value. Let's assume that in my case I'm using black and white masks where black is responsible for background and white is responsible for cars. After mask binarization all background pixels tirn into 0s and car pixels - into 1s. And here comes the bummer! When the line masks = [(mask == v) for v in self.class_values] gets executed, i get false Truevalues. In my mask all 0s are background pixels, but in class_values 0 is a car class value.

Any ideas? Should I invert colors in my mask before executing masks = [(mask == v) for v in self.class_values]?

Roopg commented 3 years ago

Hi did you figure out the fix for this problem ? Even my binark mask in inverted- white(1) for the background and 0 for the class of interest.

EtagiBI commented 3 years ago

Hi did you figure out the fix for this problem ? Even my binark mask in inverted- white(1) for the background and 0 for the class of interest.

Hello, no, I decided to switch to another repository.