from fastai.vision.all import *
import matplotlib.pyplot as plt

### Data Loader
path = Path('Spectrogram_Images')

dls = ImageDataLoaders.from_folder(
    path,
    train = '.',
     valid_pct=0.2,  # 20% of data for validation
    item_tfms = Resize(224), # Resize images to 224x224 pixels for CNN
    batch_tfms=aug_transforms(mult=1.0)  # Apply basic augmentations
)

# Show a batch of images to verify data loading
dls.show_batch(max_n=9, figsize=(8, 8))

### CNN

learn = vision_learner(dls, resnet34, metrics=accuracy)

print(learn.model)

learn.fine_tune(5)

# Plot training and validation loss curves
learn.recorder.plot_loss()

# Use the learning rate finder to identify the best learning rate
learn.lr_find()

# Plot the suggested learning rates
plt.show()