generate_partial_visual_questions.py

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

Usage:
    python generate_partial_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

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
    return icons

def create_visual_equation(equation, icons, font, icon_map):
    elements = equation.split()
    images = []
    coefficient = 1
    metadata = {}
    
    # 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")

    # Larger font for symbols and coefficients
    try:
        large_font = ImageFont.load_default().font_variant(size=30)
    except Exception:
        large_font = ImageFont.load_default()

    for idx, element in enumerate(elements):
        if element.isdigit() and idx != len(elements) - 1:
            # This is a coefficient for the next variable
            coefficient = int(element)
        elif element.isalpha():
            # This is a variable - create coefficient + icon combination
            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
            
            # Get the icon
            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)
            icon_img = selected_icon.convert("RGBA").resize((50, 50), Image.Resampling.LANCZOS)
            
            # Create coefficient text image
            coeff_img = Image.new('RGBA', (30, 50), (255, 255, 255, 255))
            draw = ImageDraw.Draw(coeff_img)
            draw.text((5, 10), str(coefficient), font=large_font, fill=(0, 0, 0))
            
            # Combine coefficient + icon horizontally
            combined_width = coeff_img.width + icon_img.width
            combined_height = max(coeff_img.height, icon_img.height)
            combined_img = Image.new('RGB', (combined_width, combined_height), (255, 255, 255))
            
            # Paste coefficient first, then icon
            if coeff_img.mode == 'RGBA':
                combined_img.paste(coeff_img, (0, 0), coeff_img)
            else:
                combined_img.paste(coeff_img, (0, 0))
            
            if icon_img.mode == 'RGBA':
                combined_img.paste(icon_img, (coeff_img.width, 0), icon_img)
            else:
                combined_img.paste(icon_img, (coeff_img.width, 0))
            
            images.append(combined_img)
            metadata[element] = (icon_map[element], coefficient)
            coefficient = 1  # Reset coefficient
        else:
            # Handle symbols (+, -, =) and constants
            if element in ['+', '-', '=', 'x']:
                symbol_img = Image.new('RGB', (30, 50), (255, 255, 255))
                draw = ImageDraw.Draw(symbol_img)
                draw.text((10, 10), element, font=large_font, fill=(0, 0, 0))
                images.append(symbol_img)
            elif idx == len(elements) - 1:  # Last element is usually the constant
                constant_img = Image.new('RGB', (50, 50), (255, 255, 255))
                draw = ImageDraw.Draw(constant_img)
                draw.text((10, 10), element, font=large_font, fill=(0, 0, 0))
                images.append(constant_img)

    # 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:
        if img.mode == 'RGBA':
            # Create a white background and paste the RGBA image on it
            temp_img = Image.new('RGB', img.size, (255, 255, 255))
            temp_img.paste(img, (0, 0), img)
            combined_image.paste(temp_img, (x_offset, 0))
        else:
            combined_image.paste(img, (x_offset, 0))
        x_offset += img.width

    return combined_image, metadata

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)

            metadata_entry = {'filename': output_path}
            for var in icon_map:
                metadata_entry[f"{var}_icon_type"] = icon_map[var]
            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:
        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 partial 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(output_folder, "metadata.csv")
    
    # Use the main function that processes all equations, not just the first one
    main(equations_file, icon_folder, output_folder, metadata_file)

if __name__ == "__main__":
    cli_main()