Custom dataset for unsupervised segmentation

for unsupervised anomaly detection at pixel level, train dataset should contain only normal data and test data should include both normal and anomaly data with their masks.


* Train--> Normal

* Test:
   - Normal
   - Not normal

* ground_truth of test data 

How can I define a custom dataset for that and process it for training and evaluation?
Can I define a dataloader for train images and a custom dataset for test (images+ground_truth)?