import os
import cv2
import numpy as np
import random
from pathlib import Path

source_dir = "litter_detection"
output_dir = "litter_detection_augmented"
class_names = ["positive", "negative"]

# change ranges
hue_shift_range = (-10, 10) 
saturation_range = (0.85, 1.15) 
brightness_range = (0.85, 1.15) 
blur_range = (1, 3)

def augment_image(image):
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV).astype(np.float32)
    
    # hue
    hue_shift = random.randint(*hue_shift_range)
    hsv[:, :, 0] = (hsv[:, :, 0] + hue_shift) % 180
    
    # sat
    sat_factor = random.uniform(*saturation_range)
    hsv[:, :, 1] = np.clip(hsv[:, :, 1] * sat_factor, 0, 255)
    
    # brightness
    brightness_factor = random.uniform(*brightness_range)
    hsv[:, :, 2] = np.clip(hsv[:, :, 2] * brightness_factor, 0, 255)
    
    augmented = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
    
    # blur
    blur_kernel = random.choice([k for k in range(*blur_range) if k % 2 == 1] + [3, 5])
    augmented = cv2.GaussianBlur(augmented, (blur_kernel, blur_kernel), 0)
    
    return augmented


for cls in class_names:
    os.makedirs(os.path.join(output_dir, cls), exist_ok=True)

total_processed = 0
for cls in class_names:
    class_dir = os.path.join(source_dir, cls)
    
    images = [f for f in os.listdir(class_dir) if os.path.isfile(os.path.join(class_dir, f))]
    images = [f for f in images if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif'))]
    
    print(f"\nProcessing {cls}: {len(images)} images...")
    
    for i, img_name in enumerate(images, 1):
        img_path = os.path.join(class_dir, img_name)
        image = cv2.imread(img_path)
        
        if image is None:
            print(f"  Failed to read {img_name}, skipping...")
            continue
        
        augmented = augment_image(image)
        
        # save image
        output_path = os.path.join(output_dir, cls, img_name)
        cv2.imwrite(output_path, augmented)
        
        total_processed += 1
        
        if i % 500 == 0:
            print(f"{i}/{len(images)}")
    
    print(f"  Completed {cls}")

print(f"\n saved to {output_dir}")
print(f"images processed: {total_processed}")