Multi class segmentation using methods used in lesson 14

I wanted to know how to modify the last lessons code to make it a multi class segmentation(pixel wise classification) when a ground truth mask, say 3 classes masks are given and a similar sized image is given.
Is there any way to get data from a custom generator, the data is present in the folders but very small no of images so I want to create data from folder when a batch size is given using a generator, and I must feed this to the ConvLearner