from PIL import Image, ImageDraw, ImageFilter
import os
import random

def create_dummy_dataset(num_images_per_class=40):
    print("Generating synthetic images...")
    
    os.makedirs('dataset/clean', exist_ok=True)
    os.makedirs('dataset/marked', exist_ok=True)
    os.makedirs('dataset/scrubbed', exist_ok=True)

    for i in range(num_images_per_class):
        r, g, b = random.randint(50, 200), random.randint(50, 200), random.randint(50, 200)
        clean_img = Image.new('RGB', (256, 256), color=(r, g, b))
        clean_img.save(f"dataset/clean/dummy_{i}.jpg")
        
        marked_img = clean_img.copy()
        draw = ImageDraw.Draw(marked_img)
        for x in range(0, 256, 32):
            draw.line([(x, 0), (x, 256)], fill=(r+20, g+20, b+20), width=2)
        for y in range(0, 256, 32):
            draw.line([(0, y), (256, y)], fill=(r+20, g+20, b+20), width=2)
        marked_img.save(f"dataset/marked/dummy_{i}.jpg")
        
        scrubbed_img = marked_img.copy().filter(ImageFilter.GaussianBlur(radius=3))
        scrubbed_img.save(f"dataset/scrubbed/dummy_{i}.jpg")

    print(f"Done! Created {num_images_per_class * 3} total images.")

if __name__ == "__main__":
    create_dummy_dataset()