search

yassouali / pytorch-segmentation

:art: Semantic segmentation models, datasets and losses implemented in PyTorch.
MIT License
1.68k stars 381 forks source link

Binary segmentation on a custom dataset #122

Closed Avigyan2001 closed 2 years ago

Avigyan2001 commented 3 years ago

I want to perform a binary segmentation on a custom dataset which contains images in 4 channels i.e. RGB-NIR. If someone could guide me as to how do I use this repo to train any model like psp-net or deeplabv3+ on my dataset. Where do I have to make all the changes in the code to perform the training and evaluation? I am lost and a help would be very much appreciated.

yassouali commented 2 years ago

Hi, @Avigyan2001

You only need to right a new dataloader (can be based on the provided ones) to load your training data, and then simply set the input_channels in the config file to 3, and you can then start training.

example:

    "arch": {
        "type": "PSPNet",
        "args": {
            "backbone": "resnet50",
            "freeze_bn": false,
            "freeze_backbone": false,
            "in_channels": 3,
        }
    },