LicensePlateDetectorDataset.py

from PIL import Image
import os
from torchvision import transforms
from torch.utils.data import Dataset
import os
import cv2
import numpy as np
import albumentations as A
from tqdm import tqdm  # tqdm provides a nice progress bar

# TODO: Load Files From Google Drive



class LicensePlateDetectorDataset(Dataset):
    ZOOM_AUGMENTATION = "zoom"
    GRAYSCALE_AUGMENTATION = "grayscale"
    FOG_AUGMENTATION = "foggy"
    RAIN_AUGMENTATION = "rainy"
    SALT_PEPPER_AUGMENTATION = "saltandpepper"
    BLUR_AUGMENTATION = "blur"

    def __init__(self, root_dir, subset: str = 'train', transform: transforms.Compose = None, augmentations = {
            "zoom": [0.5],
            "grayscale": [],
            "foggy": [],
            "rainy": [],
            "saltandpepper": [0.01, 0.01],
            "blur": [5],
        }):
        self.root_dir = os.path.join(root_dir, subset)
        self.folder_subset = subset
        self.transform = transform
        self.image_paths = []
        self.labels = []

        self.augmentations_functions = {
            LicensePlateDetectorDataset.ZOOM_AUGMENTATION: self.apply_zoom,
            LicensePlateDetectorDataset.GRAYSCALE_AUGMENTATION: self.apply_grayscale,
            LicensePlateDetectorDataset.FOG_AUGMENTATION: self.apply_fog,
            LicensePlateDetectorDataset.RAIN_AUGMENTATION: self.apply_rain,
            LicensePlateDetectorDataset.SALT_PEPPER_AUGMENTATION: self.apply_salt_and_pepper,
            LicensePlateDetectorDataset.BLUR_AUGMENTATION: self.apply_gaussian_blur,
        }

        self.augmentations = augmentations

        image_dir = os.path.join(self.root_dir, 'images')
        label_dir = os.path.join(self.root_dir, 'labels')

        # Count total number of images/files
        total_files = sum(1 for file in os.listdir(image_dir) if file.endswith(('.png', '.jpg', '.jpeg')))

        progress_bar = tqdm(total=total_files, desc='Initializing Files', unit='image')

        for img_file in os.listdir(image_dir):
            if img_file.endswith(('.png', '.jpg', '.jpeg')):
                self.image_paths.append(os.path.join(image_dir, img_file))
                progress_bar.update(1)  # Update progress bar

                print("Image:", img_file)

                label_file = os.path.join(label_dir, img_file.replace('.jpg', '.txt'))
                self.labels.append(label_file)
                print("Label:", label_file)

        progress_bar.close()

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = Image.open(img_path).convert('RGB')
        label_path = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        with open(label_path, 'r') as f:
            label = f.readline().strip()

        return image, label
    
    def process_images(self, output_folder: str, images_limit=-1):
        """Apply a given processing function to all images in the input folder and save the results to the output folder."""
        os.makedirs(output_folder, exist_ok=True)
        count = 0  # Initialize a counter for iterations
        
        # Count total number of images/files
        total_files = sum(1 for file in os.listdir(os.path.join(self.root_dir, "images")) if file.endswith(('.png', '.jpg', '.jpeg'))) if images_limit == -1 else images_limit

        progress_bar = tqdm(total=total_files, desc='Applying Augmentation to Images', unit='image')
        
        for filename in os.listdir(os.path.join(self.root_dir, "images")):
            if filename.endswith((".jpg", ".jpeg", ".png")):
                image_path = os.path.join(self.root_dir, "images", filename)
                image = cv2.imread(image_path)
                for augmentation_name, augmentation_options in self.augmentations.items():
                    processed_image = self.augmentations_functions[augmentation_name](image, *augmentation_options)
                    augmentation_folder = os.path.join(output_folder, augmentation_name)
                    os.makedirs(augmentation_folder, exist_ok=True)  # Ensure augmentation_folder exists
                    output_path = os.path.join(augmentation_folder, filename)
                    cv2.imwrite(output_path, processed_image)
                
                print("Image Processed with all augmentations: ", os.path.join(output_folder, augmentation_name, filename))
                progress_bar.update(1)  # Update progress bar

                count += 1  # Increment the counter
                if count >= images_limit and images_limit != -1: # Dont enter the loop if limit is not specified (-1 means until finish)
                    print("Reached Limit Specified", images_limit)
                    break  # Exit the loop if the limit is reached

        progress_bar.close()

    def apply_zoom(self, image: np.ndarray, scale_factor: float) -> np.ndarray:
        """Apply zoom to the image."""
        height, width = image.shape[:2]
        center = (width // 2, height // 2)
        matrix = cv2.getRotationMatrix2D(center, 0, scale_factor)
        zoomed_image = cv2.warpAffine(image, matrix, (width, height))
        return zoomed_image

    def apply_grayscale(self, image: np.ndarray) -> np.ndarray:
        """Convert image to grayscale."""
        return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    def apply_gaussian_blur(self, image: np.ndarray, kernel_size: int) -> np.ndarray:
        """Apply Gaussian blur to the image."""
        return cv2.GaussianBlur(image, (kernel_size, kernel_size), 0)

    def apply_salt_and_pepper(self, image: np.ndarray, salt_prob: float, pepper_prob: float) -> np.ndarray:
        """Apply salt and pepper noise to the image."""
        noisy_image = np.copy(image)
        row, col, _ = noisy_image.shape
        salt_pixels = np.random.rand(row, col) < salt_prob
        noisy_image[salt_pixels] = 255
        pepper_pixels = np.random.rand(row, col) < pepper_prob
        noisy_image[pepper_pixels] = 0
        return noisy_image

    def apply_fog(self, image: np.ndarray) -> np.ndarray:
        """Apply fog effect to the image."""
        transform = A.Compose([A.RandomFog(fog_coef_lower=0.3, fog_coef_upper=0.5, alpha_coef=0.3, p=1)])
        transformed = transform(image=image)
        return transformed['image']

    def apply_rain(self, image: np.ndarray) -> np.ndarray:
        """Apply rain effect to the image."""
        transform = A.Compose([A.RandomRain(brightness_coefficient=0.9, drop_width=1, blur_value=5, p=1)])
        transformed = transform(image=image)
        return transformed['image']
    
    def auto_annotate(self):
        pass
## Example Usage:
# transform = transforms.Compose([
#     transforms.Resize((192, 192)),
#     transforms.ToTensor()
# ])
# augmentations = {
    # LicensePlateDetectorDataset.ZOOM_AUGMENTATION: [15],
    # LicensePlateDetectorDataset.GRAYSCALE_AUGMENTATION: [],
    # LicensePlateDetectorDataset.FOG_AUGMENTATION: [],
    # LicensePlateDetectorDataset.RAIN_AUGMENTATION: [],
    # LicensePlateDetectorDataset.SALT_PEPPER_AUGMENTATION: [0,0],
    # LicensePlateDetectorDataset.BLUR_AUGMENTATION: [15],
# }
# lpdd = LicensePlateDetectorDataset("./License Plate YOLOv8/", subset="train", transform=transform, augmentations=augmentations)
# lpdd.process_images("./License Plate YOLOv8/augmented_images", images_limit=20) # remove images_limit or make it -1 to loop over all images