Some datasets provide a train_test_split doc, which may have image number to train/test mapping. Can we have a feature to do this?