| import os |
| import random |
| import csv |
| import urllib.request |
| from PIL import Image, ImageDraw, ImageFont |
|
|
| def _download_or_read_text(source): |
| if source.startswith("http://") or source.startswith("https://"): |
| req = urllib.request.Request(source, headers={'User-Agent': 'Mozilla/5.0'}) |
| with urllib.request.urlopen(req, timeout=10) as resp: |
| text = resp.read().decode('utf-8', errors='ignore') |
| else: |
| with open(source, 'r', encoding='utf-8', errors='ignore') as f: |
| text = f.read() |
| return text |
|
|
| def _get_words_from_text(text): |
| words = text.split() |
| words = [w.strip('.,!?"\'()[]{}') for w in words] |
| words = [w for w in words if w] |
| return words |
|
|
| def _render_text(text, bg_color=255, font_size=32): |
| try: |
| |
| font = ImageFont.truetype("arial.ttf", font_size) |
| except IOError: |
| font = ImageFont.load_default() |
| |
| |
| tmp = Image.new("L", (16, 16), 255) |
| draw = ImageDraw.Draw(tmp) |
| bbox = draw.textbbox((0, 0), text, font=font) |
| tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1] |
| |
| pad_x = random.randint(10, 20) |
| pad_y = random.randint(5, 15) |
| |
| cw = max(32, tw + pad_x * 2) |
| ch = max(32, th + pad_y * 2) |
| |
| image = Image.new("RGB", (cw, ch), (bg_color, bg_color, bg_color)) |
| draw = ImageDraw.Draw(image) |
| |
| x = pad_x |
| y = pad_y |
| |
| draw.text((x, y), text, fill=(0, 0, 0), font=font) |
| |
| |
| angle = random.uniform(-2, 2) |
| image = image.rotate(angle, resample=Image.Resampling.BICUBIC, expand=True, fillcolor=(bg_color,bg_color,bg_color)) |
| |
| return image |
|
|
| def generate_silly_training_set(source, count, output_dir="silly_train"): |
| """ |
| Generates images containing individual words extracted from 'source'. |
| 'Beta' utility for quick training on custom datasets. |
| Returns the path to the generated labels.csv file. |
| """ |
| os.makedirs(output_dir, exist_ok=True) |
| text = _download_or_read_text(source) |
| words = _get_words_from_text(text) |
| |
| if not words: |
| raise ValueError("Source has no valid words.") |
| |
| labels_file = os.path.join(output_dir, "labels.csv") |
| |
| with open(labels_file, "w", encoding="utf-8", newline="") as f: |
| writer = csv.writer(f) |
| writer.writerow(["image_path", "label"]) |
| |
| for i in range(count): |
| word = random.choice(words) |
| img = _render_text(word) |
| |
| img_name = f"train_{i:06d}.jpg" |
| img_path = os.path.join(output_dir, img_name) |
| img.save(img_path) |
| |
| writer.writerow([img_path, word]) |
| |
| if i % 100 == 0: |
| print(f"Generated {i}/{count} training samples...", end="\r") |
| |
| print(f"\nSaved {count} training samples to {labels_file}") |
| return labels_file |
|
|
| def generate_silly_testing_set(source, count, output_dir="silly_test"): |
| """ |
| Generates images containing sentences of 2-6 words extracted from 'source'. |
| 'Beta' utility for testing inference on continuous text. |
| Returns the path to the generated labels.csv file. |
| """ |
| os.makedirs(output_dir, exist_ok=True) |
| text = _download_or_read_text(source) |
| words = _get_words_from_text(text) |
| |
| if not words: |
| raise ValueError("Source has no valid words.") |
| |
| labels_file = os.path.join(output_dir, "labels.csv") |
| |
| with open(labels_file, "w", encoding="utf-8", newline="") as f: |
| writer = csv.writer(f) |
| writer.writerow(["image_path", "label"]) |
| |
| for i in range(count): |
| sentence_len = random.randint(2, 6) |
| sentence_words = [random.choice(words) for _ in range(sentence_len)] |
| sentence = " ".join(sentence_words) |
| |
| img = _render_text(sentence) |
| |
| img_name = f"test_{i:06d}.jpg" |
| img_path = os.path.join(output_dir, img_name) |
| img.save(img_path) |
| |
| writer.writerow([img_path, sentence]) |
| |
| if i % 100 == 0: |
| print(f"Generated {i}/{count} testing samples...", end="\r") |
| |
| print(f"\nSaved {count} testing samples to {labels_file}") |
| return labels_file |
|
|
