Spaces:
Sleeping
Sleeping
| import argparse | |
| import os | |
| import sys | |
| from pathlib import Path | |
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) | |
| from src.data_prep.canonicalize import run_canonicalization | |
| from src.data_prep.collisions import run_collision_removal | |
| from src.data_prep.compile import run_compilation | |
| from src.data_prep.normalize import run_normalization | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Сборка датасета Graph2Eq") | |
| parser.add_argument( | |
| "--raw_data", type=str, default="data", help="Папка с сырыми чанками" | |
| ) | |
| parser.add_argument( | |
| "--output_dir", | |
| type=str, | |
| default="data_cleared", | |
| help="Папка для .pt тензоров", | |
| ) | |
| args = parser.parse_args() | |
| # Папки для промежуточных этапов | |
| raw_dir = Path(args.raw_data) | |
| step1_dir = Path("data/tmp_1_canonical") | |
| step2_dir = Path("data/tmp_2_normalized") | |
| step3_dir = Path("data/tmp_3_no_collisions") | |
| final_dir = Path(args.output_dir) | |
| print("\n" + "=" * 40) | |
| print("ШАГ 1: Канонизация выражений (cos -> sin)") | |
| print("=" * 40) | |
| run_canonicalization(raw_dir, step1_dir) | |
| print("\n" + "=" * 40) | |
| print("ШАГ 2: Геометрическая чистка и BBox нормализация") | |
| print("=" * 40) | |
| run_normalization(step1_dir, step2_dir) | |
| print("\n" + "=" * 40) | |
| print("ШАГ 3: Удаление визуальных коллизий (GPU)") | |
| print("=" * 40) | |
| ckpt_path = Path("data/gpu_checkpoint.pt") | |
| run_collision_removal(step2_dir, step3_dir, ckpt_path) | |
| print("\n" + "=" * 40) | |
| print("ШАГ 4: Компиляция PyTorch тензоров") | |
| print("=" * 40) | |
| run_compilation(step3_dir, final_dir) | |
| print("\n Датасет готов!") | |
| if __name__ == "__main__": | |
| main() | |