Data augmentation on Car Collision Dataset using Marimo

augment_dataset.py

from concurrent.futures import ThreadPoolExecutor, as_completed

# Try to import tqdm for a nice live progress bar; fallback to a basic counter if not installed
try:
    from tqdm import tqdm

    HAS_TQDM = True
except ImportError:
    HAS_TQDM = False

@mo.persistent_cache
def augment_dataset():
    """
    Augment the train set from the dataset.
    """

    original_images = get_original_images(train_only=True)
    total_images = len(original_images)
    print(f"Found {total_images} original images. Starting parallel processing...")

    with ThreadPoolExecutor(max_workers=None) as executor:
        # Submit all image tasks to the pool
        futures = {
            executor.submit(augment_single_image, path): path for path in original_images
        }

        # Monitor execution progress
        if HAS_TQDM:
            # Displays a beautiful live updating progress bar
            for _ in tqdm(as_completed(futures), total=total_images, desc="Augmenting Images"):
                pass
        else:
            completed_count = 0
            for future in as_completed(futures):
                completed_count += 1
                if completed_count % 500 == 0 or completed_count == total_images:
                    print(
                        f"Progress: {completed_count}/{total_images} original images processed."
                    )

    print("All images successfully augmented in parallel!")

augment_dataset()

augment_single_image.py

@mo.persistent_cache
def augment_single_image(image_path):
    image = cv2.imread(image_path)
    if image is None:
        return False

    folder_path = os.path.dirname(image_path)
    filename, ext = os.path.splitext(os.path.basename(image_path))

    # 1. Apply Exposure Variations
    for suffix, factor in exposure_tasks.items():
        exposure_image = adjust_exposure(image, factor)
        cv2.imwrite(
            filename=os.path.join(folder_path, f"{filename}{suffix}{ext}"),
            img=exposure_image,
        )

    # 2. Apply Noise Variations
    for suffix, percentage in noise_tasks.items():
        noise_image = add_salt_and_pepper_noise(image, percentage)
        cv2.imwrite(
            filename=os.path.join(folder_path, f"{filename}{suffix}{ext}"),
            img=noise_image,
        )

    # 3. Apply Horizontal Flip
    flip_image = apply_horizontal_flip(image)
    cv2.imwrite(
        filename=os.path.join(folder_path, f"{filename}_horizontal_flip{ext}"),
        img=flip_image,
    )

    # 4. Apply Fixed Zoom Motion Blur
    z_blur_image = apply_zoom_motion_blur(image)
    cv2.imwrite(
        filename=os.path.join(folder_path, f"{filename}_motion_blur{ext}"),
        img=z_blur_image,
    )

    # 5. Apply Gaussian Blur
    g_blur_image = apply_gaussian_blur(image)
    cv2.imwrite(
        filename=os.path.join(folder_path, f"{filename}_gaussian_blur{ext}"),
        img=g_blur_image,
    )

    return True

description.md

Augment Data

Data augmentation should include:

• Darken by 5%, 10%, 15%, 20% (under-exposure)
• Brighen by 5%, 10%, 15%, 20% (over-exposure)
• Horizontal Flip
• Motion blur
• Gaussian blur
• Salt-and-pepper Noise

functions.py

import cv2
from pathlib import Path

# Safely construct folder paths using os.path.join
dataset_splits = {
    "train": os.path.join(base_path, "dataset", "train"),
    "validation": os.path.join(base_path, "dataset", "valid"),
    "test": os.path.join(base_path, "dataset", "test"),
}

def adjust_exposure(image, factor):
    """
    Darkens or brightens the image by a given scale factor.
    Using float32 prevents pixel value overflow/underflow wrapping.
    """

    # Steps:
    # 1. Convert to float
    # 2. Scale pixel values
    # 3. Clip between 0-255
    rescaled = np.clip(image.astype(np.float32) * factor, 0, 255)
    return rescaled.astype(np.uint8)

def add_salt_and_pepper_noise(image, percentage):
    """Adds a specific percentage of random black and white noise pixels."""
    noisy_image = image.copy()
    h, w, _ = image.shape
    random_matrix = np.random.rand(h, w)

    # salt (white)
    salt_mask = random_matrix < (percentage / 2)
    noisy_image[salt_mask] = [255, 255, 255]

    # pepper (black)
    pepper_mask = (random_matrix >= (percentage / 2)) & (random_matrix < percentage)
    noisy_image[pepper_mask] = [0, 0, 0]

    return noisy_image

def apply_zoom_motion_blur(image, num_frames=10, zoom_range=0.05):
    """
    Applies a zoom motion blur to an image.
    Simulates a zoom motion blur by scaling the image around its center point multiple times and averaging the frames together.

    Args:
        image (numpy.ndarray): The input image (H, W, C).
        num_frames (int): Number of zoomed layers to stack.
        zoom_range (float): Magnitude of the zoom effect.
    """
    height, width = image.shape[:2]
    center_x = width / 2
    center_y = height / 2

    # Create an empty float32 array to safely accumulate image matrices without clipping
    blur_accumulator = np.zeros_like(image, dtype=np.float32)

    # Generate 10 evenly spaced zoom scale factors from 1.0 (original)
    scales = np.linspace(1.0, 1.0 + zoom_range, num_frames)

    for scale in scales:
        # Get the transformation matrix for scaling around the center (0 degree rotation)
        M = cv2.getRotationMatrix2D((center_x, center_y), 0, scale)
        # Apply the transformation to the image
        warped = cv2.warpAffine(image, M, (width, height), flags=cv2.INTER_LINEAR)
        # Add the frame into our accumulator
        blur_accumulator += warped

    # Divide by total frames to get the mean effect, then convert back to 8-bit image
    zoom_blurred_image = blur_accumulator / num_frames
    return zoom_blurred_image.astype(np.uint8)

def apply_gaussian_blur(image, kernel_size=7):
    """
    Applies standard Gaussian Blur. Kernel size must be positive and odd.
    """
    return cv2.GaussianBlur(image, (kernel_size, kernel_size), sigmaX=0)

def apply_horizontal_flip(image):
    """
    Flips the image horizontally (axis 1).
    """
    return cv2.flip(image, flipCode=1)

get_original_images.py

@mo.persistent_cache
def get_original_images(train_only: bool = True):
    original_images = []
    for split_name, split_path in dataset_splits.items():
        if train_only and split_name != "train":
            continue

        image_count_per_split = 0
        print(f"Gathering original files for {split_name} set")

        if not os.path.exists(split_path):
            print(f"[warning] Path {split_path} does not exist. Skipping...")
            continue

        all_suffixes = (
            list(exposure_tasks.keys())
            + list(noise_tasks.keys())
            + ["_horizontal_flip", "_motion_blur", "_gaussian_blur"]
        )
        for root, dirs, files in os.walk(split_path):
            for file in files:
                if file.lower().endswith(valid_extensions):
                    # Ensure the file doesn't already contain one of our augmentation suffixes
                    if not any(suffix in file for suffix in all_suffixes):
                        original_images.append(os.path.join(root, file))
                        image_count_per_split += 1

        print(f"Found {image_count_per_split} original images in {split_name} set.")
    return original_images

suffixes.py

exposure_tasks = {
    "_under_exposure_5p": 0.95,
    "_under_exposure_10p": 0.90,
    "_under_exposure_15p": 0.85,
    "_under_exposure_20p": 0.80,
    "_over_exposure_5p": 1.05,
    "_over_exposure_10p": 1.10,
    "_over_exposure_15p": 1.15,
    "_over_exposure_20p": 1.20,
}

noise_tasks = {
    "_noise_5p": 0.05,
    "_noise_10p": 0.10,
    "_noise_15p": 0.15,
    "_noise_20p": 0.20,
}

valid_extensions = (".jpg", ".jpeg", ".png", ".bmp")