File size: 7,474 Bytes
bd5a354 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 |
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 |