generate_data.py - master - nacha/pseado

generate_data.py
52 строки2 KB
alex
Add synthetic data generation script
07 авг 2026, 08:38
07 авг 2026, 08:38fb1ee2c
"""
Скрипт для генерации тестовых данных (синтетических).
Создает структуру директорий и заглушки для датасета.
"""
import os
import json
import random

def main():
    base_dir = "data/drawings"
    images_dir = os.path.join(base_dir, "images")
    annotations_dir = os.path.join(base_dir, "annotations")

    os.makedirs(images_dir, exist_ok=True)
    os.makedirs(annotations_dir, exist_ok=True)

    num_samples = 10
    
    # Имитация имен файлов (можно использовать UUID или номера чертежей)
    for i in range(num_samples):
        filename = f"draft_{i:03d}.jpg"
        # Создаем пустой файл изображения (чтобы загрузка не падала)
        full_img_path = os.path.join(images_dir, filename)
        with open(full_img_path, "wb") as f:
            # Простой минимальный JPG-файл (1x1 red pixel) - не валидный, но лучше бы PIL. 
            # Здесь просто создаем пустой для теста директорий.
            # В реальности: PIL Image.new('RGB', (512, 512)).save(f)
            f.write(b"") 

        # Создаем JSON аннотацию
        num_nodes = random.randint(5, 20)
        # Генерируем случайные связи для GNN
        edges = []
        for _ in range(random.randint(1, 10)):
            u = random.randint(0, num_nodes - 1)
            v = random.randint(0, num_nodes - 1)
            edges.append([u, v])
            
        annotation = {
            "score": random.uniform(50, 100),
            "labels": [1 if random.random() > 0.7 else 0 for _ in range(5)],
            "num_nodes": num_nodes,
            "graph": edges
        }
        
        with open(os.path.join(annotations_dir, filename.replace('.jpg', '.json')), 'w') as f:
            json.dump(annotation, f, indent=2)
            
    print(f"Generated {num_samples} synthetic samples in {base_dir}")

if __name__ == "__main__":
    main()