For best practice:

I would advise to separate train/val/test earlier, possibly even at the file system level.
I know nothing about your dataset, but you have checked that just splitting at a image level is the right thing to do (in medical imaging, you typically want to have all images of a given patient on the same side of the split and the like)?
After the above two, it is natural to have separate datasets for train/val/test. At that point, having different transforms is easy. 


Best regards
Thomas