SAM2 Refactoring + Speedup Attempt

[Kitchi]

Trying to speedup the data augmentation + transformation portions of SAM2 training by offloading them on the GPU.

The plan is to use the kornia library (pyTorch compatible) to do this. Eventually wondering if we can get to something like NVIDIA DALI to bypass the CPU entirely and speed up IO off disk to the GPU directly, for further speedups. But that is more of a stretch goal.

[preshanth]

@Kitchi Thanks for the classes. I have used them to setup a faster GPU based augmentation and also realized that the polarizations were leading to worse validation scores. So I have turned that off and now both training and validations are in a more regularized mode with significant weight decay. So we can go to 25 epochs now and are able to keep the H100 completely hot. The batch rate is still a bit low but at this point it is for the lack of GPU resources not for what we have done. I have extended the sample to about 12K samples for training and 1k for validation. I have also added randomization in terms of how much RFI is there too so SAM does not learn the total amount. This is rounding out quite well. If the validation script on real data shows the improvement we think we can call it done. I want to merge this branch and finish clean up on all the developments. I am going to create a Main PR issue so we can cleanup and do the other stuff such as docs and code cleanup.