generate_dataset.py

"""
Procedurally generate a synthetic texture dataset for DCGAN training.
Generates 5 texture types: wood, marble, fabric, brick, noise.
Each type produces 200 images → 1000 total, saved to data/textures/<type>/.
"""
import os
import random
import numpy as np
from PIL import Image, ImageFilter, ImageDraw

IMG_SIZE = 64
NUM_PER_CLASS = 200
OUTPUT_DIR = os.path.join("data", "textures")

TEXTURE_TYPES = ["wood", "marble", "fabric", "brick", "noise"]


def generate_wood(size=IMG_SIZE):
    """Simulate wood grain using sine waves + noise."""
    arr = np.zeros((size, size, 3), dtype=np.uint8)
    freq = random.uniform(3, 8)
    phase = random.uniform(0, np.pi * 2)
    for y in range(size):
        for x in range(size):
            grain = np.sin(freq * y / size * np.pi * 2 + phase + random.gauss(0, 0.05))
            r = int(np.clip(120 + 60 * grain + random.randint(-10, 10), 80, 200))
            g = int(np.clip(70 + 30 * grain + random.randint(-10, 10), 40, 130))
            b = int(np.clip(30 + 10 * grain + random.randint(-5, 5), 10, 70))
            arr[y, x] = [r, g, b]
    img = Image.fromarray(arr)
    img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
    return img


def generate_marble(size=IMG_SIZE):
    """Simulate marble veins using Perlin-like noise."""
    arr = np.zeros((size, size, 3), dtype=np.uint8)
    scale = random.uniform(0.05, 0.15)
    for y in range(size):
        for x in range(size):
            val = np.sin(x * scale + y * scale * 0.5 + random.gauss(0, 0.3)) * 0.5 + 0.5
            base = int(200 + 40 * val)
            r = int(np.clip(base + random.randint(-5, 5), 150, 255))
            g = int(np.clip(base - 10 + random.randint(-5, 5), 140, 245))
            b = int(np.clip(base - 5 + random.randint(-5, 5), 145, 250))
            arr[y, x] = [r, g, b]
    img = Image.fromarray(arr)
    img = img.filter(ImageFilter.GaussianBlur(radius=0.3))
    return img


def generate_fabric(size=IMG_SIZE):
    """Simulate woven fabric using a grid pattern."""
    arr = np.zeros((size, size, 3), dtype=np.uint8)
    thread_size = random.randint(3, 6)
    hue_r = random.randint(50, 200)
    hue_g = random.randint(50, 200)
    hue_b = random.randint(50, 200)
    for y in range(size):
        for x in range(size):
            is_warp = (x // thread_size) % 2 == 0
            is_weft = (y // thread_size) % 2 == 0
            if is_warp and not is_weft:
                factor = 1.2
            elif not is_warp and is_weft:
                factor = 0.8
            else:
                factor = 1.0
            r = int(np.clip(hue_r * factor + random.randint(-5, 5), 0, 255))
            g = int(np.clip(hue_g * factor + random.randint(-5, 5), 0, 255))
            b = int(np.clip(hue_b * factor + random.randint(-5, 5), 0, 255))
            arr[y, x] = [r, g, b]
    return Image.fromarray(arr)


def generate_brick(size=IMG_SIZE):
    """Simulate brick wall pattern."""
    img = Image.new("RGB", (size, size), color=(180, 80, 50))
    draw = ImageDraw.Draw(img)
    brick_h = random.randint(8, 12)
    brick_w = random.randint(16, 24)
    mortar_color = (200, 190, 180)
    for row in range(size // brick_h + 1):
        y = row * brick_h
        offset = (brick_w // 2) if row % 2 else 0
        # Horizontal mortar line
        draw.line([(0, y), (size, y)], fill=mortar_color, width=1)
        # Vertical mortar lines
        for col in range(-1, size // brick_w + 2):
            x = col * brick_w + offset
            draw.line([(x, y), (x, y + brick_h)], fill=mortar_color, width=1)
    # Add slight noise
    arr = np.array(img)
    noise = np.random.randint(-15, 15, arr.shape, dtype=np.int16)
    arr = np.clip(arr.astype(np.int16) + noise, 0, 255).astype(np.uint8)
    return Image.fromarray(arr)


def generate_noise(size=IMG_SIZE):
    """Generate colorful noise texture."""
    mode = random.choice(["perlin_like", "color_bands", "static"])
    if mode == "static":
        arr = np.random.randint(0, 255, (size, size, 3), dtype=np.uint8)
    elif mode == "color_bands":
        arr = np.zeros((size, size, 3), dtype=np.uint8)
        for y in range(size):
            r = int((np.sin(y * 0.3) * 0.5 + 0.5) * 255)
            g = int((np.cos(y * 0.2) * 0.5 + 0.5) * 255)
            b = int((np.sin(y * 0.5 + 1) * 0.5 + 0.5) * 255)
            arr[y, :] = [r, g, b]
        noise = np.random.randint(-30, 30, arr.shape, dtype=np.int16)
        arr = np.clip(arr.astype(np.int16) + noise, 0, 255).astype(np.uint8)
    else:
        arr = np.zeros((size, size, 3), dtype=np.uint8)
        scale = random.uniform(0.1, 0.3)
        for y in range(size):
            for x in range(size):
                v = int((np.sin(x * scale) * np.cos(y * scale) * 0.5 + 0.5) * 255)
                arr[y, x] = [v, int(v * 0.7), int(v * 0.4)]
    img = Image.fromarray(arr)
    img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
    return img


GENERATORS = {
    "wood": generate_wood,
    "marble": generate_marble,
    "fabric": generate_fabric,
    "brick": generate_brick,
    "noise": generate_noise,
}


def main():
    total = 0
    for texture_type in TEXTURE_TYPES:
        out_dir = os.path.join(OUTPUT_DIR, texture_type)
        os.makedirs(out_dir, exist_ok=True)
        gen_fn = GENERATORS[texture_type]
        print(f"Generating {NUM_PER_CLASS} '{texture_type}' textures...")
        for i in range(NUM_PER_CLASS):
            img = gen_fn(IMG_SIZE)
            img.save(os.path.join(out_dir, f"{texture_type}_{i:04d}.png"))
            total += 1
        print(f"  OK Saved {NUM_PER_CLASS} images to {out_dir}")

    print(f"\n Dataset generation complete! Total images: {total}")
    print(f"   Saved to: {os.path.abspath(OUTPUT_DIR)}")


if __name__ == "__main__":
    main()