src/operators/main_augmentation.operator.py

import os
import cv2
import numpy as np
from typing import Dict, List, Tuple, Optional, Set
from tqdm import tqdm
from datetime import datetime
import random
import shutil
import logging
from pathlib import Path
import glob

logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(levelname)s - %(message)s')

def get_image_count(folder_path: str) -> int:
    """Get exact count of valid image files in folder."""
    pattern = os.path.join(folder_path, "*.[pjJ][npNP][gG]*")
    return len(glob.glob(pattern))

def load_image(image_path: str) -> Optional[np.ndarray]:
    """Load and preprocess image."""
    try:
        image = cv2.imread(image_path)
        if image is None:
            return None
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        if np.mean(gray) > 127:
            gray = 255 - gray
        return gray
    except Exception:
        return None

def apply_augmentation(image: np.ndarray, aug_type: str, value: float) -> np.ndarray:
    """Apply augmentation transformation."""
    if aug_type == "ROT":
        h, w = image.shape[:2]
        center = (w // 2, h // 2)
        matrix = cv2.getRotationMatrix2D(center, value, 1.0)
        return cv2.warpAffine(image, matrix, (w, h),
                            borderMode=cv2.BORDER_CONSTANT,
                            borderValue=0)
    else:  # SCL
        h, w = image.shape[:2]
        new_h = int(h * value)
        new_w = int(w * value)
        scaled = cv2.resize(image, (new_w, new_h))
        result = np.zeros((h, w), dtype=np.uint8)
        y_offset = (h - new_h) // 2
        x_offset = (w - new_w) // 2
        if y_offset >= 0 and x_offset >= 0:
            result[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = scaled
        else:
            y_start = abs(min(y_offset, 0))
            x_start = abs(min(x_offset, 0))
            result = scaled[y_start:y_start+h, x_start:x_start+w]
        return result

def generate_augmentations(needed: int) -> List[Tuple[str, float]]:
    """Generate augmentation parameters ensuring enough variations."""
    augmentations = []
    
    # More granular rotation angles
    angles = [-4.0, -3.5, -3.0, -2.5, -2.0, -1.5, -1.0, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0]
    
    # More granular scaling factors
    scales = [0.995, 0.9965, 0.998, 0.9995, 1.0005, 1.002, 1.0035, 1.005]
    
    # Calculate repetitions needed
    total_variations = len(angles) + len(scales)
    repetitions = (needed // total_variations) + 1
    
    for _ in range(repetitions):
        for angle in angles:
            augmentations.append(("ROT", angle))
        for scale in scales:
            augmentations.append(("SCL", scale))
    
    random.shuffle(augmentations)
    return augmentations[:needed]

def process_character_folder(char_folder: str, target_count: int = 20000) -> None:
    """Process single character folder ensuring exact target count."""
    images_path = os.path.join(char_folder, 'images')
    if not os.path.exists(images_path):
        logging.error(f"Images path does not exist: {images_path}")
        return

    current_count = get_image_count(images_path)
    if current_count == target_count:
        logging.info(f"Folder {os.path.basename(char_folder)} already has {target_count} images.")
        return

    if current_count > target_count:
        all_images = glob.glob(os.path.join(images_path, "*.[pjJ][npNP][gG]*"))
        to_remove = random.sample(all_images, current_count - target_count)
        
        for img_path in tqdm(to_remove, desc="Removing excess images"):
            os.remove(img_path)
        return

    # Load base images
    base_images = []
    all_images = glob.glob(os.path.join(images_path, "*.[pjJ][npNP][gG]*"))
    
    for img_path in tqdm(all_images, desc="Loading base images"):
        img = load_image(img_path)
        if img is not None:
            base_images.append((img, os.path.splitext(os.path.basename(img_path))[0]))

    if not base_images:
        logging.error(f"No valid base images found in {images_path}")
        return

    needed = target_count - current_count
    augmentations = generate_augmentations(needed)
    unique_identifier = datetime.now().strftime("%H%M%S")
    generated = 0
    
    with tqdm(total=needed, desc=f"Generating images for {os.path.basename(char_folder)}") as pbar:
        for idx, (aug_type, value) in enumerate(augmentations):
            if generated >= needed:
                break
                
            base_img, base_name = random.choice(base_images)
            aug_img = apply_augmentation(base_img, aug_type, value)
            
            filename = f"{base_name}_{aug_type}_{int(abs(value*1000)):04d}_{unique_identifier}_{idx:06d}.png"
            save_path = os.path.join(images_path, filename)
            
            if cv2.imwrite(save_path, aug_img):
                generated += 1
                pbar.update(1)

    final_count = get_image_count(images_path)
    if final_count < target_count:
        remaining = target_count - final_count
        logging.info(f"Generating remaining {remaining} images...")
        process_character_folder(char_folder, target_count)
    elif final_count > target_count:
        excess = final_count - target_count
        logging.info(f"Removing {excess} excess images...")
        all_images = glob.glob(os.path.join(images_path, "*.[pjJ][npNP][gG]*"))
        to_remove = random.sample(all_images, excess)
        for img_path in to_remove:
            os.remove(img_path)

def balance_dataset(base_path: str) -> None:
    """Balance entire dataset to exactly 20,000 images per character."""
    categories = ['digits', 'uppercase']  # Only digits and uppercase
    target_count = 20000
    
    for category in tqdm(categories, desc="Processing categories"):
        category_path = os.path.join(base_path, category)
        if not os.path.exists(category_path):
            logging.error(f"Category path does not exist: {category_path}")
            continue

        char_folders = sorted([f for f in os.listdir(category_path) 
                             if os.path.isdir(os.path.join(category_path, f))])
        
        for char_folder in tqdm(char_folders, desc=f"Processing {category}", leave=False):
            full_path = os.path.join(category_path, char_folder)
            process_character_folder(full_path, target_count)

def main():
    base_path = "/Users/roammer/Documents/Github/handwriting/data/processed"
    logging.info("Starting dataset balancing...")
    balance_dataset(base_path)
    logging.info("Dataset balancing completed!")

if __name__ == "__main__":
    main()