data_preparation.py

import os
import pandas as pd
from sklearn.model_selection import train_test_split
from tensorflow.keras.preprocessing.image import ImageDataGenerator

def get_generators(data_dir, batch_size=16, img_size=(224, 224)):
    # Generate data paths with labels
    filepaths = []
    labels = []

    # Assuming directories are named "cancerous" and "non-cancerous"
    for folder in ['cancerous', 'non-cancerous']:
        folder_path = os.path.join(data_dir, folder)
        if os.path.isdir(folder_path):  # Ensure it's a valid folder
            for root, _, files in os.walk(folder_path):
                for file in files:
                    # Check for valid image extensions
                    if file.endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff')):
                        file_path = os.path.join(root, file)
                        filepaths.append(file_path)
                        labels.append(folder)  # Use folder name as label
        else:
            print(f"Directory not found: {folder_path}")

    # Debugging: Print stats about the dataset
    print(f"Total files: {len(filepaths)}")
    print(f"Labels: {set(labels)}")  # Should show {'cancerous', 'non-cancerous'}

    # Create a DataFrame
    df = pd.DataFrame({'filepaths': filepaths, 'labels': labels})
    df['labels'] = df['labels'].map({'non-cancerous': 0, 'cancerous': 1})  # Map labels to 0/1

    if df.empty:
        raise ValueError("Dataset is empty! Please check the dataset path and files.")

    # Split the data into train, validation, and test datasets
    train_df, dummy_df = train_test_split(df, train_size=0.8, shuffle=True, random_state=123)
    valid_df, test_df = train_test_split(dummy_df, train_size=0.6, shuffle=True, random_state=123)

    # Initialize ImageDataGenerators
    tr_gen = ImageDataGenerator(rescale=1.0/255)  # Normalize pixel values
    ts_gen = ImageDataGenerator(rescale=1.0/255)

    # Create data generators
    train_gen = tr_gen.flow_from_dataframe(
        train_df, x_col='filepaths', y_col='labels', target_size=img_size,
        class_mode='raw', color_mode='rgb', shuffle=True, batch_size=batch_size
    )

    valid_gen = ts_gen.flow_from_dataframe(
        valid_df, x_col='filepaths', y_col='labels', target_size=img_size,
        class_mode='raw', color_mode='rgb', shuffle=True, batch_size=batch_size
    )

    test_gen = ts_gen.flow_from_dataframe(
        test_df, x_col='filepaths', y_col='labels', target_size=img_size,
        class_mode='raw', color_mode='rgb', shuffle=False, batch_size=batch_size
    )
    
    return train_gen, valid_gen, test_gen