src/preprocessing.py

import os
from typing import Tuple

import tensorflow as tf

def delete_corrupted_image(dataset_path: str, categories: Tuple[str]) -> int:
    num_skipped = 0
    
    for folder_name in categories:
        folder_path = os.path.join(dataset_path, folder_name)
        for fname in os.listdir(folder_path):
            fpath = os.path.join(folder_path, fname)
            try:
                fobj = open(fpath, 'rb')
                is_jfif = tf.compat.as_bytes("JFIF") in fobj.peek(10)
            finally:
                fobj.close()

            if not is_jfif:
                num_skipped += 1
                os.remove(fpath)
            
    return num_skipped

def get_data_augmentation() -> tf.keras.Sequential:
    return tf.keras.Sequential([
        tf.keras.layers.RandomFlip('horizontal'),
        tf.keras.layers.RandomRotation(0.1)
    ])