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