Dataset structure and training script

[HongLiuuuuu]

Thanks for this repo and this impressive work! I am trying to train on my own dataset. Could you please provide the dataset structure and training script?

[joihn]

Here is what I've reverse engineered for the original dataloader

path_to_your_dataset/
   img1.jpg
   img2.jpg
   img3.jpg
   img1_mask.png
   img2_mask.png
   img3_mask.png

Make sure you add you path path_to_your_dataset in the Trainer class of the train.py (take inspiration from config/image_gen_train.yaml)

Make sure you use you own config system (unless you are in a jupyter notebook, vanilla one has some issue #3)

then

python -m accelerate.commands.launch train.py

Ps: original dataloader has quite some hardcoded stuff for their specific dataset, you might be better off rewriting yours

[HongLiuuuuu]

Thanks for your help!

[HongLiuuuuu]

I am trying to write the dataloader for my own dataset. However, it is hard for me to understand the returned mask and label in dataset_masks.py without the 'labels.csv' and 'labels.pickle' files. Could you please show some examples of the two files or explain?

[joihn]

here is a hint for your custom dataloder

from torch.utils.data import Dataset,DataLoader
class DatasetCustom(Dataset):

        """
        output must be (C x H x W)
        img is a torch tensor of 3x512x512, in the range [0.0, 1.0]
        mask is a torch tensor of 1x512x512
          the idx value in the mask should be :
          0: class unknown (mandatory class)
          1: your_first_class
          2 : your_second_class

          with p=0.5, you should feed an empty mask (0 everywhere)(unlabeled data)

        """
        # your code here
        
        return img, label

[bhosalems]

@joihn or @HongLiuuuuu are you able to write the dataset class for training dm.py models? I believe that you need data for each label like 'stroma' ...., and then the mask generation model will be used to generate masks that can combine these classes together in the conditional mask based on the contextual promot like "signet ring cell carcinoma" to generate an arbitrary large image. They will use models trained in dm.py by applying the generation per class and taking the union of all generations as shown in Fig. 2b right part of the paper. @marcoaversa can you confirm this or give us some examples of how the datasets can be arranged so that minimal change is required on the training part of the code?