Language/Musnad_Final_Collection/generate_data.py at main · AdamAI-Systems/Language · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import os
import random
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageFilter

# Constants
FONTS_DIR = os.path.join(os.path.dirname(__file__), "Fonts")
CLASSES_FILE = os.path.join(os.path.dirname(__file__), "classes.txt")
OUTPUT_DIR = "dataset"
IMG_SIZE = 224

# Load Musnad classes
with open(CLASSES_FILE, "r", encoding="utf-8") as f:
    classes = [line.strip() for line in f if line.strip()]

# Load available fonts
fonts = [os.path.join(FONTS_DIR, f) for f in os.listdir(FONTS_DIR) if f.lower().endswith(('.ttf', '.otf'))]

def generate_sample(img_id, split="train"):
    # Create background with random noise
    bg_color = random.randint(220, 255)
    img_array = np.full((IMG_SIZE, IMG_SIZE, 3), bg_color, dtype=np.uint8)
    noise = np.random.randint(-15, 15, (IMG_SIZE, IMG_SIZE, 3), dtype=np.int16)
    img_array = np.clip(img_array.astype(np.int16) + noise, 0, 255).astype(np.uint8)

    img = Image.fromarray(img_array)
    draw = ImageDraw.Draw(img)
    num_objs = random.randint(1, 4)
    labels = []

    for _ in range(num_objs):
        char_idx = random.randint(0, len(classes) - 1)
        char = classes[char_idx]
        font_path = random.choice(fonts)
        font_size = random.randint(35, 55)

        try:
            font = ImageFont.truetype(font_path, font_size)
        except: continue

        bbox = draw.textbbox((0, 0), char, font=font)
        w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]

        # Draw on temp layer for rotation
        char_img = Image.new("RGBA", (w+20, h+20), (0,0,0,0))
        char_draw = ImageDraw.Draw(char_img)
        color = random.randint(0, 60)
        char_draw.text((10, 10), char, font=font, fill=(color, color, color, 255))

        # Rotate
        angle = random.randint(-20, 20)
        rotated = char_img.rotate(angle, expand=True, resample=Image.BICUBIC)
        rw, rh = rotated.size

        x = random.randint(0, IMG_SIZE - rw)
        y = random.randint(0, IMG_SIZE - rh)
        img.paste(rotated, (x, y), rotated)

        # YOLO labels (relative centers and sizes)
        x_c = (x + rw/2) / IMG_SIZE
        y_c = (y + rh/2) / IMG_SIZE
        rel_w = rw / IMG_SIZE
        rel_h = rh / IMG_SIZE
        labels.append(f"{char_idx} {x_c:.6f} {y_c:.6f} {rel_w:.6f} {rel_h:.6f}")

    # Final touch: subtle blur
    if random.random() > 0.5:
        img = img.filter(ImageFilter.GaussianBlur(radius=random.uniform(0.1, 0.4)))

    # Save
    img.convert("RGB").save(os.path.join(OUTPUT_DIR, "images", split, f"{img_id:06d}.jpg"))
    with open(os.path.join(OUTPUT_DIR, "labels", split, f"{img_id:06d}.txt"), "w", encoding="utf-8") as f:
        f.write("\n".join(labels))

def main():
    # Setup directories
    for s in ["train", "val", "test"]:
        os.makedirs(os.path.join(OUTPUT_DIR, "images", s), exist_ok=True)
        os.makedirs(os.path.join(OUTPUT_DIR, "labels", s), exist_ok=True)

    counts = {"train": 2000, "val": 300, "test": 100}
    for split, count in counts.items():
        print(f"Generating {count} samples for {split}...")
        for i in range(count):
            generate_sample(i, split)
    print("Dataset generation complete!")

if __name__ == "__main__":
    main()