ChartPipeline / scripts /process_layouts.py
Ray1ee01's picture
Upload folder using huggingface_hub
58e6885 verified
Raw
History Blame Contribute Delete
2.08 kB
import os
import json
import shutil
from pathlib import Path
def process_layouts(input_dir: str, output_dir: str):
"""
Process SVG and JSON files based on their layout information.
Args:
input_dir: Directory containing input SVG and JSON files
output_dir: Base directory for output files
"""
# Create output directory if it doesn't exist
os.makedirs(output_dir, exist_ok=True)
# Get all JSON files
json_files = list(Path(input_dir).glob("*.json"))
for json_file in json_files:
# Get corresponding SVG file
svg_file = json_file.with_suffix('.svg')
if not svg_file.exists():
print(f"Warning: No corresponding SVG file found for {json_file}")
continue
try:
# Read JSON file
with open(json_file, 'r', encoding='utf-8') as f:
layout_info = json.load(f)
# Get layout information
title_to_chart = layout_info.get('title_to_chart', 'unknown')
image_to_chart = layout_info.get('image_to_chart', 'unknown')
# Create new file name with layout info as prefix
new_filename = f"{title_to_chart}_{image_to_chart}_{svg_file.stem}.svg"
# Copy SVG file to new location
output_svg_path = os.path.join(output_dir, new_filename)
shutil.copy2(svg_file, output_svg_path)
print(f"Processed {svg_file.name} -> {new_filename}")
except Exception as e:
print(f"Error processing {json_file}: {e}")
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Process SVG files based on layout information")
parser.add_argument("--input", type=str, required=True, help="Input directory containing SVG and JSON files")
parser.add_argument("--output", type=str, required=True, help="Output base directory")
args = parser.parse_args()
process_layouts(args.input, args.output)