A couple of questions about an UNET CNN implementation in Tensorflow

So I am a beginner in this field, for a summer project, I took datasets of chest CT scans and lesions and aimed to segment them, and to diagnose specific ILDs(interstitial lung diseases) based on them. My chest segmentation model runs very well(acc>98%), but its the lesion segmentation part(which I hope will diagnose specific diseases) which is giving me problems. So my model is a multi-class model, with 17 classes(same as labels and features too right?) such as ground_glass, fibrosis, bronchial_wall_thickening and so on,and the way it works is that if the input has a specific set of these features, a specific disease can be diagnosed. 17 classes seem too much for my system(32 gb DDR4, RTX 2060 mobile), and the code crashes during the train-test split part. The code runs well if I do not read the full dataset(which contains 1724 train and 431 test images, all 512x512x1), but then I get confused which classes are being processed, and how significant are the parameter values. How should I proceed to run my model without my IDE crashing due to RAM overload, and will Colab pro do the trick? Also what can I optimize in my code, will resizing the images to 128x128x1 help? And if yes how do I proceed with that?

P.S: Here my code(the dataset is not uploaded haha but my thought process would be better understood there)

P.Sx2: Also posted this on the DeepLearning sub, my apologies if you had to read this twice.