""" Data Augmentation Pipeline för Zoomer-submissions Ökar dataset-storlek 10-20x via syntetisk variation """ import os import random import numpy as np from pathlib import Path from typing import List, Dict, Tuple, Optional from dataclasses import dataclass import json from PIL import Image, ImageEnhance, ImageFilter, ImageOps import cv2 @dataclass class AugmentationConfig: """Konfiguration för data augmentation""" # Geometrisk rotation_range: Tuple[float, float] = (-15, 15) scale_range: Tuple[float, float] = (0.8, 1.2) shear_range: Tuple[float, float] = (-5, 5) flip_horizontal_prob: float = 0.5 flip_vertical_prob: float = 0.0 # Fotometrisk brightness_range: Tuple[float, float] = (0.7, 1.3) contrast_range: Tuple[float, float] = (0.7, 1.3) saturation_range: Tuple[float, float] = (0.5, 1.5) hue_shift_range: Tuple[float, float] = (-10, 10) # Brus & kvalitet gaussian_noise_prob: float = 0.3 gaussian_noise_std: Tuple[float, float] = (5, 15) jpeg_compression_prob: float = 0.3 jpeg_quality_range: Tuple[int, int] = (60, 95) blur_prob: float = 0.2 blur_radius_range: Tuple[float, float] = (0.5, 2.0) # Väder & ljus rain_prob: float = 0.1 fog_prob: float = 0.1 shadow_prob: float = 0.15 overexposure_prob: float = 0.1 underexposure_prob: float = 0.1 # Multiplicerare target_multiplier: int = 10 class DataAugmenter: """ Augmenterar bilder för att simulera verkliga Zoomer-submissions Simulerar: - Olika telefoner (iPhone, Samsung, Xiaomi) - Olika ljusförhållanden (dagsljus, skymning, mörker) - Olika väder (regn, dimma, sol) - Olika vinklar och avstånd - Olika bildkvalitet (kompression, brus) """ def __init__(self, config: Optional[AugmentationConfig] = None): self.config = config or AugmentationConfig() self.augmentation_stats = { "total_generated": 0, "by_type": {} } def augment_image( self, image: Image.Image, bboxes: Optional[List[List[float]]] = None, labels: Optional[List[int]] = None ) -> Tuple[Image.Image, Optional[List[List[float]]], Optional[List[int]]]: """ Applicera augmentation på en bild Args: image: PIL Image bboxes: YOLO-format bboxes [x_center, y_center, width, height] labels: Klass-labels Returns: Augmenterad bild, uppdaterade bboxes, labels """ img = image.copy() new_bboxes = bboxes.copy() if bboxes else None new_labels = labels.copy() if labels else None # 1. Geometrisk transformation img, new_bboxes = self._apply_geometric(img, new_bboxes) # 2. Fotometrisk transformation img = self._apply_photometric(img) # 3. Brus & kvalitet img = self._apply_noise_and_quality(img) # 4. Väder & ljus img = self._apply_weather(img) return img, new_bboxes, new_labels def _apply_geometric( self, img: Image.Image, bboxes: Optional[List[List[float]]] ) -> Tuple[Image.Image, Optional[List[List[float]]]]: """Applicera geometrisk transformation""" # Rotation angle = random.uniform(*self.config.rotation_range) img = img.rotate(angle, resample=Image.BILINEAR, expand=False) # Skalning scale = random.uniform(*self.config.scale_range) w, h = img.size new_w, new_h = int(w * scale), int(h * scale) img = img.resize((new_w, new_h), Image.LANCZOS) # Crop tillbaka till original storlek if new_w > w or new_h > h: left = (new_w - w) // 2 top = (new_h - h) // 2 img = img.crop((left, top, left + w, top + h)) else: # Pad om nödvändigt padded = Image.new('RGB', (w, h), (128, 128, 128)) left = (w - new_w) // 2 top = (h - new_h) // 2 padded.paste(img, (left, top)) img = padded # Horisontell flip if random.random() < self.config.flip_horizontal_prob: img = ImageOps.mirror(img) if bboxes: new_bboxes = [] for bbox in bboxes: x_center, y_center, width, height = bbox new_bboxes.append([1.0 - x_center, y_center, width, height]) bboxes = new_bboxes # Vertikal flip (sällan för infrastruktur) if random.random() < self.config.flip_vertical_prob: img = ImageOps.flip(img) if bboxes: new_bboxes = [] for bbox in bboxes: x_center, y_center, width, height = bbox new_bboxes.append([x_center, 1.0 - y_center, width, height]) bboxes = new_bboxes return img, bboxes def _apply_photometric(self, img: Image.Image) -> Image.Image: """Applicera fotometrisk transformation""" # Ljusstyrka brightness = random.uniform(*self.config.brightness_range) enhancer = ImageEnhance.Brightness(img) img = enhancer.enhance(brightness) # Kontrast contrast = random.uniform(*self.config.contrast_range) enhancer = ImageEnhance.Contrast(img) img = enhancer.enhance(contrast) # Mättnad saturation = random.uniform(*self.config.saturation_range) enhancer = ImageEnhance.Color(img) img = enhancer.enhance(saturation) # Skärpa sharpness = random.uniform(0.8, 1.2) enhancer = ImageEnhance.Sharpness(img) img = enhancer.enhance(sharpness) return img def _apply_noise_and_quality(self, img: Image.Image) -> Image.Image: """Applicera brus och kvalitetsförsämring""" # Gaussiskt brus if random.random() < self.config.gaussian_noise_prob: img_array = np.array(img) noise_std = random.uniform(*self.config.gaussian_noise_std) noise = np.random.normal(0, noise_std, img_array.shape).astype(np.int16) img_array = np.clip(img_array.astype(np.int16) + noise, 0, 255).astype(np.uint8) img = Image.fromarray(img_array) # JPEG-kompression if random.random() < self.config.jpeg_compression_prob: quality = random.randint(*self.config.jpeg_quality_range) import io buffer = io.BytesIO() img.save(buffer, format='JPEG', quality=quality) buffer.seek(0) img = Image.open(buffer) # Blur if random.random() < self.config.blur_prob: radius = random.uniform(*self.config.blur_radius_range) img = img.filter(ImageFilter.GaussianBlur(radius=radius)) return img def _apply_weather(self, img: Image.Image) -> Image.Image: """Simulera vädereffekter""" # Regn if random.random() < self.config.rain_prob: img = self._add_rain(img) # Dimma if random.random() < self.config.fog_prob: img = self._add_fog(img) # Skuggor if random.random() < self.config.shadow_prob: img = self._add_shadow(img) # Överexponering if random.random() < self.config.overexposure_prob: enhancer = ImageEnhance.Brightness(img) img = enhancer.enhance(1.5) enhancer = ImageEnhance.Contrast(img) img = enhancer.enhance(0.7) # Underexponering if random.random() < self.config.underexposure_prob: enhancer = ImageEnhance.Brightness(img) img = enhancer.enhance(0.5) enhancer = ImageEnhance.Contrast(img) img = enhancer.enhance(1.3) return img def _add_rain(self, img: Image.Image) -> Image.Image: """Lägg till regneffekt""" img_array = np.array(img) h, w = img_array.shape[:2] # Skapa regndroppar num_drops = random.randint(500, 2000) for _ in range(num_drops): x = random.randint(0, w - 1) y = random.randint(0, h - 1) length = random.randint(5, 20) intensity = random.randint(150, 255) if y + length < h: img_array[y:y+length, x] = intensity return Image.fromarray(img_array) def _add_fog(self, img: Image.Image) -> Image.Image: """Lägg till dimma""" img_array = np.array(img).astype(np.float32) h, w = img_array.shape[:2] # Skapa gradient fog_intensity = random.uniform(0.1, 0.3) fog = np.ones_like(img_array) * 255 * fog_intensity # Applicera fog img_array = img_array * (1 - fog_intensity) + fog return Image.fromarray(np.clip(img_array, 0, 255).astype(np.uint8)) def _add_shadow(self, img: Image.Image) -> Image.Image: """Lägg till skuggor""" img_array = np.array(img).astype(np.float32) h, w = img_array.shape[:2] # Skapa slumpmässig skugga shadow_mask = np.ones((h, w), dtype=np.float32) # Lägg till flera skuggrektanglar num_shadows = random.randint(1, 3) for _ in range(num_shadows): x1 = random.randint(0, w // 2) y1 = random.randint(0, h // 2) x2 = random.randint(x1 + 50, w) y2 = random.randint(y1 + 50, h) intensity = random.uniform(0.3, 0.7) shadow_mask[y1:y2, x1:x2] = intensity # Applicera skugga for c in range(3): img_array[:, :, c] *= shadow_mask return Image.fromarray(np.clip(img_array, 0, 255).astype(np.uint8)) def augment_dataset( self, input_dir: str, output_dir: str, annotations_file: Optional[str] = None ) -> Dict: """ Augmentera hela dataset Args: input_dir: Input bildkatalog output_dir: Output bildkatalog annotations_file: COCO/YOLO annotations JSON Returns: Statistik över augmentation """ input_path = Path(input_dir) output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) # Ladda annotations annotations = {} if annotations_file and os.path.exists(annotations_file): with open(annotations_file) as f: annotations = json.load(f) # Hitta alla bilder image_files = list(input_path.glob("*.jpg")) + list(input_path.glob("*.png")) stats = { "original_images": len(image_files), "generated_images": 0, "augmentations_applied": {} } print(f"Augmenting {len(image_files)} images with {self.config.target_multiplier}x multiplier...") for img_file in image_files: # Ladda original img = Image.open(img_file).convert('RGB') # Spara original original_out = output_path / f"{img_file.stem}_orig.jpg" img.save(original_out, quality=95) # Generera augmenterade varianter for i in range(self.config.target_multiplier): aug_img, _, _ = self.augment_image(img) # Spara aug_filename = f"{img_file.stem}_aug_{i:03d}.jpg" aug_out = output_path / aug_filename aug_out.save(aug_img, quality=90) stats["generated_images"] += 1 # Uppdatera statistik aug_types = self._get_applied_augmentations() for aug_type in aug_types: stats["augmentations_applied"][aug_type] = \ stats["augmentations_applied"].get(aug_type, 0) + 1 print(f"\nAugmentation complete!") print(f" Original images: {stats['original_images']}") print(f" Generated images: {stats['generated_images']}") print(f" Total dataset size: {stats['original_images'] + stats['generated_images']}") return stats def _get_applied_augmentations(self) -> List[str]: """Hämta vilka augmentationer som applicerades""" # Förenklad version - returnera slumpmässig subset all_augs = [ "rotation", "scale", "flip", "brightness", "contrast", "saturation", "noise", "jpeg_compression", "blur", "rain", "fog", "shadow", "overexposure", "underexposure" ] # Välj 3-7 slumpmässiga augmentationer num_augs = random.randint(3, 7) return random.sample(all_augs, num_augs) def create_augmented_training_set( real_images_dir: str, output_dir: str, target_size: int = 10000 ) -> Dict: """ Skapa augmenterat träningsdataset från riktiga bilder Args: real_images_dir: Katalog med riktiga Zoomer-submissions output_dir: Output-katalog target_size: Målstorlek på dataset Returns: Dataset-statistik """ real_path = Path(real_images_dir) # Räkna riktiga bilder real_images = list(real_path.glob("*.jpg")) + list(real_path.glob("*.png")) num_real = len(real_images) if num_real == 0: raise ValueError(f"No images found in {real_images_dir}") # Beräkna multiplier multiplier = max(1, target_size // num_real - 1) print(f"Real images: {num_real}") print(f"Target size: {target_size}") print(f"Multiplier: {multiplier}x") # Skapa augmenter config = AugmentationConfig(target_multiplier=multiplier) augmenter = DataAugmenter(config) # Augmentera stats = augmenter.augment_dataset(real_images_dir, output_dir) return { "real_images": num_real, "target_size": target_size, "multiplier": multiplier, "actual_size": stats["original_images"] + stats["generated_images"], "augmentation_stats": stats } # Exempel if __name__ == "__main__": # Skapa augmenterat dataset result = create_augmented_training_set( real_images_dir="/tmp/quixzoom_training/images", output_dir="/tmp/quixzoom_augmented", target_size=5000 ) print("\n" + "=" * 50) print("AUGMENTATION COMPLETE") print("=" * 50) print(f"Real images: {result['real_images']}") print(f"Target size: {result['target_size']}") print(f"Actual size: {result['actual_size']}")