generate_visual_questions.py

"""
Script to generate visual questions for math-eval dataset.

Usage:
    python generate_visual_questions.py --input_dir <input_dir> --output_dir <output_dir> --icon_dir <icon_dir> [other flags]
"""
import os
from PIL import Image, ImageDraw, ImageFont
from tqdm import tqdm
import random
import csv
import argparse

# Function to load icons from the dataset
def load_icons(icon_folder):
    icons = {}
    for root, _, files in os.walk(icon_folder):
        icon_name = os.path.basename(root)
        if icon_name not in icons:
            icons[icon_name] = []
        for filename in files:
            if filename.endswith(".png"):
                icons[icon_name].append(Image.open(os.path.join(root, filename)))
                break
    print(icons.keys())
    return icons

# Function to create visual equations
def create_visual_equation(equation, icons, font, icon_map):
    elements = equation.split()
    images = []
    coefficient = 1
    metadata = {}
    
    # Get list of available icon types
    available_icon_types = [k for k, v in icons.items() if v]
    if len(available_icon_types) < 2:
        raise ValueError("Not enough icon types available")

    # Create a larger font
    large_font = ImageFont.load_default().font_variant(size=30)  # Adjust size value as needed

    for idx, element in enumerate(elements):
        if element.isdigit() and idx != len(elements) - 1:
            coefficient = int(element)
        elif element.isalpha():
            # If element not in map, assign a random icon type different from existing ones
            if element not in icon_map:
                used_types = set(icon_map.values())
                available_types = [t for t in available_icon_types if t not in used_types]
                if not available_types:
                    raise ValueError("Not enough unique icon types available")
                icon_type = random.choice(available_types)
                icon_map[element] = icon_type
            
            icon_variations = icons.get(icon_map[element], [])
            if not icon_variations:
                raise ValueError(f"No icons available for type {icon_map[element]}")
                
            selected_icon = random.choice(icon_variations)
            for _ in range(coefficient):
                images.append(selected_icon.resize((50, 50), Image.Resampling.LANCZOS))
            metadata[element] = (icon_map[element], coefficient)
            coefficient = 1
        else:
            if element in ['+', '-', '=', 'x']:
                symbol_image = Image.new('RGBA', (50, 50), (255, 255, 255, 0))
                draw = ImageDraw.Draw(symbol_image)
                draw.text((15, 5), element, font=large_font, fill=(0, 0, 0))
                images.append(symbol_image)
            if idx == len(elements) - 1:
                constant_image = Image.new('RGBA', (50, 50), (255, 255, 255, 0))
                draw = ImageDraw.Draw(constant_image)
                draw.text((15, 5), element, font=large_font, fill=(0, 0, 0))
                images.append(constant_image)

    # Combine images side by side
    widths, heights = zip(*(i.size for i in images))
    total_width = sum(widths)
    max_height = max(heights)

    combined_image = Image.new('RGB', (total_width, max_height), color=(255, 255, 255))

    x_offset = 0
    for img in images:
        combined_image.paste(img, (x_offset, 0), img if img.mode == 'RGBA' else None)
        x_offset += img.width

    return combined_image, metadata

# Main function to process the equations file and generate images
def main(equations_file, icon_folder, output_folder, metadata_file):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    icons = load_icons(icon_folder)
    font = ImageFont.load_default()

    with open(equations_file, 'r') as file:
        lines = file.readlines()

    metadata_list = []

    for i, line in tqdm(enumerate(lines), desc="Processing equations", total=len(lines)):
        try:
            # Split the line to get equations and variable values
            full_line = line.strip().split('<sep>')
            equations_set = full_line[0].strip()
            variable_values = {}
            if len(full_line) > 1:
                # Parse variable values from input (assuming format like "a=2,b=3")
                value_pairs = full_line[1].strip().split(',')
                for pair in value_pairs:
                    if '=' in pair:
                        var, val = pair.strip().split('=')
                        variable_values[var.strip()] = int(val.strip())

            equations = equations_set.split(',')
            icon_map = {}
            equation_images = []
            
            # Process equations to generate images
            for equation in equations:
                img, eq_metadata = create_visual_equation(equation.strip(), icons, font, icon_map)
                equation_images.append(img)

            # Combine images vertically
            widths, heights = zip(*(img.size for img in equation_images))
            max_width = max(widths)
            total_height = sum(heights) + (20 * (len(equation_images) - 1))

            combined_image = Image.new('RGB', (max_width, total_height), color=(255, 255, 255))

            y_offset = 0
            for img in equation_images:
                combined_image.paste(img, (0, y_offset))
                y_offset += img.height + 20

            output_path = os.path.join(output_folder, f"equation_set_{i + 1}.png")
            combined_image.save(output_path)

            # Create metadata entry with both icon mapping and variable values
            metadata_entry = {
                'filename': output_path,
            }
            
            # Add icon mapping information
            for var in icon_map:
                metadata_entry[f"{var}_icon_type"] = icon_map[var]
            
            # Add variable values
            for var, value in variable_values.items():
                metadata_entry[f"{var}_value"] = value

            metadata_list.append(metadata_entry)

        except ValueError as e:
            print(f"Skipping equation set {i + 1}: {str(e)}")
            continue

    # Write metadata to CSV file
    with open(metadata_file, 'w', newline='') as csvfile:
        # Get all possible field names
        fieldnames = ['filename']
        for d in metadata_list:
            fieldnames.extend(k for k in d.keys() if k not in fieldnames)
        
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        for data in metadata_list:
            writer.writerow(data)

def cli_main():
    parser = argparse.ArgumentParser(description="Generate visual questions for math-eval dataset.")
    parser.add_argument('--equations_file', type=str, required=True, help='Path to the equations file')
    parser.add_argument('--icon_dir', type=str, required=True, help='Directory containing icon images')
    parser.add_argument('--output_dir', type=str, required=True, help='Output directory for generated questions')
    args = parser.parse_args()

    equations_file = args.equations_file
    icon_folder = args.icon_dir
    output_folder = args.output_dir
    metadata_file = os.path.join(args.output_dir, "metadata.csv")
    
    # Call the actual main function with proper parameters
    main(equations_file, icon_folder, output_folder, metadata_file)

if __name__ == "__main__":
    cli_main()