| 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 |
|
|
| |
|
|
|
|
|
|
| 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') |
|
|
| |
| 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) |
|
|
| 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 |
| |
| |
| 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) |
| 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) |
|
|
| count += 1 |
| if count >= images_limit and images_limit != -1: |
| print("Reached Limit Specified", images_limit) |
| break |
|
|
| 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 |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
