generator.py - master - tuper/Coursovaya

generator.py
184 строки8 KB
tuper
create: demo_text.py, generator.py, generator_text.py, demo.py
29 май 2026, 16:27
29 май 2026, 16:27524d2b7
import os
import json
import random
import numpy as np
import cv2
import urllib.request
from datetime import datetime
from PIL import Image, ImageDraw, ImageFont
import albumentations as A

# Инициализация Faker
try:
    from faker import Faker
    fake = Faker('ru_RU')
except ImportError:
    print("Установите faker: pip install faker")
    exit()

class AugmentedLetterGenerator:
    def __init__(self, width=1240, height=1754):
        self.width = width
        self.height = height
        
        # Автоматическое скачивание кириллических шрифтов, чтобы не было "квадратиков"
        self._ensure_fonts_exist()
        
        self.font_regular = ImageFont.truetype("Roboto-Regular.ttf", 24)
        self.font_bold = ImageFont.truetype("Roboto-Bold.ttf", 26)
        self.font_small = ImageFont.truetype("Roboto-Regular.ttf", 18)

        # Обновленный пайплайн Albumentations (совместим с новыми версиями)
        self.transform = A.Compose([
            # Поворот (без устаревших параметров)
            A.Rotate(limit=3, p=0.7),
            
            # Искажение перспективы
            A.Perspective(scale=(0.02, 0.05), p=0.5),
            
            # Шум (используем MultiplicativeNoise, он стабильнее между версиями)
            A.MultiplicativeNoise(multiplier=(0.9, 1.1), p=0.6),
            
            # Размытие
            A.MotionBlur(blur_limit=5, p=0.3),
            
            # Контраст и яркость
            A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.8),
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))

    def _ensure_fonts_exist(self):
        """Скачивает шрифты, если их нет в текущей директории"""
        fonts = {
            "Roboto-Regular.ttf": "https://github.com/googlefonts/roboto/raw/main/src/hinted/Roboto-Regular.ttf",
            "Roboto-Bold.ttf": "https://github.com/googlefonts/roboto/raw/main/src/hinted/Roboto-Bold.ttf"
        }
        for font_name, url in fonts.items():
            if not os.path.exists(font_name):
                print(f"Скачивание шрифта {font_name} для поддержки кириллицы...")
                urllib.request.urlretrieve(url, font_name)

    def generate_document_data(self):
        corr_types = [
            f"{fake.company()}\n{fake.city()}, {fake.street_address()}\nТел: {fake.phone_number()}",
            f"Министерство цифрового развития\nЧелябинской области\nДепартамент управления делами"
        ]
        return {
            "correspondent": random.choice(corr_types),
            "letter_date": fake.date(pattern="%d.%m.%Y"),
            "letter_number": f"исх-{random.randint(100, 999)}/{random.randint(20, 26)}",
            "addressee": f"{fake.company()}\nДиректору по развитию\n{fake.name()}",
            "signatory": f"Директор\n\n\n_________________ / {fake.name()} /"
        }

    def draw_text_and_get_bbox(self, draw, text, x, y, font):
        lines = text.split('\n')
        current_y = y
        max_w = 0
        
        for line in lines:
            if not line:
                current_y += 25
                continue
            
            # Поддержка разных версий Pillow
            try:
                left, top, right, bottom = draw.textbbox((x, current_y), line, font=font)
                w = right - left
                h = bottom - top
            except AttributeError:
                w, h = draw.textsize(line, font=font)
                
            draw.text((x, current_y), line, font=font, fill=0)
            if w > max_w: max_w = w
            current_y += h + 8
            
        return [x, y, x + max_w, current_y]

    def create_and_augment_letter(self, output_img_path, output_json_path):
        image = Image.new("L", (self.width, self.height), 255)
        draw = ImageDraw.Draw(image)
        data = self.generate_document_data()
        
        bboxes = []
        labels = []

        # 1. Корреспондент
        bbox_corr = self.draw_text_and_get_bbox(draw, data["correspondent"], 100, 100, self.font_small)
        bboxes.append(bbox_corr)
        labels.append("correspondent")
        
        draw.line([(100, bbox_corr[3] + 20), (self.width - 100, bbox_corr[3] + 20)], fill=128, width=2)
        
        # 2. Адресат
        bbox_addr = self.draw_text_and_get_bbox(draw, data["addressee"], self.width - 500, bbox_corr[3] + 50, self.font_regular)
        bboxes.append(bbox_addr)
        labels.append("addressee")
        
        # 3. Дата и Номер
        meta_y = bbox_addr[1]
        bbox_date = self.draw_text_and_get_bbox(draw, f"Дата: {data['letter_date']}", 100, meta_y, self.font_regular)
        bboxes.append(bbox_date)
        labels.append("letter_date")
        
        bbox_num = self.draw_text_and_get_bbox(draw, f"№ {data['letter_number']}", 100, meta_y + 40, self.font_regular)
        bboxes.append(bbox_num)
        labels.append("letter_number")
        
        # Тело письма (чтобы скан не казался пустым)
        body_y = max(bbox_num[3], bbox_addr[3]) + 100
        draw.text((self.width//2 - 150, body_y), "ИНФОРМАЦИОННОЕ ПИСЬМО", font=self.font_bold, fill=0)
        
        fake_body = (
            "Настоящим уведомляем Вас о проведении плановых работ.\n"
            "В период с 01.06.2026 по 03.06.2026 возможны задержки в обработке\n"
            "входящей корреспонденции.\n\n"
            "Надеемся на плодотворное сотрудничество."
        )
        self.draw_text_and_get_bbox(draw, fake_body, 100, body_y + 80, self.font_regular)
        
        # 4. Подписант
        bbox_sig = self.draw_text_and_get_bbox(draw, data["signatory"], 100, self.height - 300, self.font_regular)
        bboxes.append(bbox_sig)
        labels.append("signatory")

        # АУГМЕНТАЦИЯ
        image_np = np.array(image)
        transformed = self.transform(image=image_np, bboxes=bboxes, class_labels=labels)
        
        transformed_image = transformed['image']
        transformed_bboxes = transformed['bboxes']
        transformed_labels = transformed['class_labels']

        # СОХРАНЕНИЕ
        final_img = Image.fromarray(transformed_image)
        final_img.save(output_img_path)
        
        final_annotations = {
            "image_path": output_img_path,
            "bboxes": {}
        }
        for bbox, label in zip(transformed_bboxes, transformed_labels):
            final_annotations["bboxes"][label] = [int(coord) for coord in bbox]
            
        with open(output_json_path, 'w', encoding='utf-8') as f:
            json.dump(final_annotations, f, ensure_ascii=False, indent=4)

if __name__ == "__main__":
    os.makedirs("dataset/images/train", exist_ok=True)
    os.makedirs("dataset/labels/train", exist_ok=True)
    
    generator = AugmentedLetterGenerator()
    
    num_samples = 100
    print(f"Генерация {num_samples} аугментированных писем...")
    
    for i in range(num_samples):
        img_p = f"dataset/images/train/doc_{i:04d}.png"
        json_p = f"dataset/labels/train/doc_{i:04d}.json"
        
        generator.create_and_augment_letter(img_p, json_p)
        
        if (i + 1) % 10 == 0:
            print(f"Готово: {i + 1}/{num_samples}")
            
    print("\nДатасет успешно сгенерирован!")