DerradjAdel/Lifeline_labs_AI

1import io2import numpy as np3from PIL import Image4from tensorflow.keras.preprocessing.image import ImageDataGenerator5 6# Adjust target size to match your 3-class model input (224×224)7TARGET_SIZE = (299, 299)8 9# Optional augmentation setup10augmenter = ImageDataGenerator(11    rescale=1./255,12    rotation_range=10,13    width_shift_range=0.01,14    height_shift_range=0.01,15    zoom_range=0.01,16    shear_range=0.01,17    brightness_range=[0.9, 1.5],  # simulate X-ray exposure differences18    fill_mode='constant' 19)20 21def load_and_resize(image_bytes, target_size=TARGET_SIZE):22    img = Image.open(io.BytesIO(image_bytes)).convert("L")23    return img.resize(target_size)24 25def augment_image(img: Image.Image):26    """27    Apply a random augmentation to the given PIL Image.28 29    Returns a new PIL Image (grayscale).30    """31    # Convert PIL image to array (height, width)32    arr = np.array(img)33    # Add channel dimension: (height, width) -> (height, width, 1)34    arr = np.expand_dims(arr, axis=-1)35    # Add batch dimension: (height, width, 1) -> (1, height, width, 1)36    batch = np.expand_dims(arr, axis=0)37    # Apply augmentation: output shape is (1, height, width, 1)38    transformed = next(augmenter.flow(batch, batch_size=1))39    # Remove batch dim: (1, h, w, 1) -> (h, w, 1)40    arr_aug = transformed[0]41    # Convert back to 2D grayscale by squeezing channel42    arr_aug = np.clip(arr_aug * 255, 0, 255).astype('uint8').squeeze(-1)43    return Image.fromarray(arr_aug, mode='L')44 45 46def to_array(img: Image.Image):47    """48    Convert PIL Image to normalized NumPy array suitable for model input.49    Output shape: (1, height, width, 1)50    """51    arr = np.array(img).astype('float32') / 255.0  # (h, w)52    # Add channel and batch dims53    arr = np.expand_dims(arr, axis=-1)  # (h, w, 1)54    return np.expand_dims(arr, axis=0)  # (1, h, w, 1)55 56def preprocess_for_model(image_bytes, augment=False):57    img = load_and_resize(image_bytes)58    if augment:59        img = augment_image(img)60    return to_array(img)61