SCMayS/hydata

1import os2import re3import glob4import json5import argparse6import random7import uuid8from tqdm import tqdm9from pathlib import Path10from collections import defaultdict11 12def parse_ground_truth(name):13    """Extract ground truth rotation axis and angle from filename or folder name"""14    # Remove file extension if present15    basename = name.split(".")[0] if "." in name else name16    17    parts = basename.split("_")18    if len(parts) >= 4:  # figXXXX_XXX_axis_angle19        rotation_axis = parts[-2]  # Second to last element is axis20        rotation_angle = int(parts[-1])  # Last element is angle21        22        # Convert negative angles to 0-360 range23        if rotation_angle < 0:24            rotation_angle += 36025            26        return rotation_axis, rotation_angle27    28    print(f"Warning: Could not parse name: {basename}")29    return None, None30 31def load_examples(example_dir, generation_mode):32    """Load example images from the example directory"""33    if generation_mode == "combined":34        # Load all single PNG files from the example directory35        files = glob.glob(os.path.join(example_dir, "*.png"))36        print(f"Found {len(files)} combined example images in {example_dir}")37        return files38    else:  # separate mode39        # Find all folders in the example directory40        folders = [f for f in glob.glob(os.path.join(example_dir, "*")) if os.path.isdir(f)]41        # Filter folders that contain both _ini.png and _rot.png files42        valid_folders = []43        for folder in folders:44            folder_name = os.path.basename(folder)45            ini_file = os.path.join(folder, f"{folder_name}_ini.png")46            rot_file = os.path.join(folder, f"{folder_name}_rot.png")47            if os.path.exists(ini_file) and os.path.exists(rot_file):48                valid_folders.append(folder)49        50        print(f"Found {len(valid_folders)} example folder pairs in {example_dir}")51        return valid_folders52 53def organize_examples(examples, generation_mode):54    """Organize examples by rotation axis and angle"""55    organized = defaultdict(list)56    for example in examples:57        basename = os.path.basename(example)58        if generation_mode == "combined":59            basename = basename.split(".")[0]60            61        axis, angle = parse_ground_truth(basename)62        if axis is None or angle is None:63            continue64            65        key = (axis, angle)66        organized[key].append(example)67    68    # Print statistics69    print("\nDistribution of examples by axis-angle:")70    for key, examples_list in organized.items():71        print(f"  {key[0]}-axis, {key[1]} degrees: {len(examples_list)} examples")72    73    return dict(organized)74 75def select_examples(organized_examples, test_axis, possible_angles, max_examples=2):76    """Select examples for the test case, with appropriate randomization"""77    examples = []78    79    # Organize examples by angle for this axis80    examples_by_angle = {}81    for (axis, angle), example_list in organized_examples.items():82        if axis == test_axis and angle in possible_angles:83            if angle not in examples_by_angle:84                examples_by_angle[angle] = []85            examples_by_angle[angle].extend([(example, angle) for example in example_list])86    87    # If no examples found, return empty list88    if not examples_by_angle:89        print(f"Warning: No examples found for rotation around {test_axis}-axis")90        return []91    92    # If max_examples is less than the number of possible angles,93    # randomly select max_examples angles and pick one example from each94    if max_examples < len(possible_angles):95        # Get available angles that have examples96        available_angles = [angle for angle in possible_angles if angle in examples_by_angle and examples_by_angle[angle]]97        98        # Randomly select angles to use99        if available_angles:100            selected_angles = random.sample(available_angles, min(max_examples, len(available_angles)))101            102            # Select one example from each selected angle103            for angle in selected_angles:104                selected_example = random.choice(examples_by_angle[angle])105                examples.append(selected_example)106    else:107        # Try to select one example from each angle108        for angle in possible_angles:109            if angle in examples_by_angle and examples_by_angle[angle]:110                selected_example = random.choice(examples_by_angle[angle])111                examples.append(selected_example)112                113                # If we have enough examples, stop114                if len(examples) >= max_examples:115                    break116    117        # If we don't have enough examples, fill with random examples from any angle118        if len(examples) < max_examples and examples_by_angle:119            all_examples = []120            for angle_examples in examples_by_angle.values():121                all_examples.extend(angle_examples)122            123            while len(examples) < max_examples and all_examples:124                selected_example = random.choice(all_examples)125                # Avoid duplicates126                all_examples.remove(selected_example)127                if selected_example not in examples:128                    examples.append(selected_example)129    130    return examples131 132def construct_prompt_with_examples(axis, possible_angles, examples=None, difficulty="easy", generation_mode="combined"):133    """Create prompt for the VLM with an in-context example"""134    # Generate list of all possible rotation angles based on angle increment135    '''possible_angles = []136    current_angle = 0 + angle_increment137    while current_angle < 360:138        possible_angles.append(current_angle)139        current_angle += angle_increment'''140    141    # Common instructions for both modes142    coordinate_system_templates = [143        # Version 1 (Slightly more concise)144        (145            "The 3D Cartesian coordinate setup is:\n"146            "- x-axis: Runs horizontally, positive to the right.\n"147            "- y-axis: Runs vertically, positive going up.\n"148            "- z-axis: Runs perpendicular to the screen, positive towards the viewer.\n"149            "- The origin (0,0,0) is positioned at the geometric center of the 3D object mesh.\n\n"150            "To visualize rotations, imagine looking along the positive direction of the rotation axis."151        ),152 153        # Version 2 (Focus on orientation and direction)154        (155            "We are using a 3D Cartesian system oriented as follows:\n"156            "- The positive x-direction points right.\n"157            "- The positive y-direction points up.\n"158            "- The positive z-direction points out of the screen towards you.\n"159            "- The coordinate system's origin (0,0,0) coincides with the geometric center of the 3D object.\n\n"160            "When discussing axis rotations, the viewpoint is assumed to be looking along the axis's positive direction."161        ),162 163        # Version 3 (More explicit about axes)164        (165            "This work employs a 3D Cartesian coordinate frame where:\n"166            "- The horizontal axis (X) increases positively to the right.\n"167            "- The vertical axis (Y) increases positively upwards.\n"168            "- The depth axis (Z) increases positively coming out towards the observer.\n"169            "- The reference point (0,0,0) is located at the object's geometric center.\n\n"170            "For rotations, the perspective is looking down the positive axis towards the origin."171        ),172        173        # Version 4 (Technical description)174        (175            "The coordinate system for this task is defined as:\n"176            "- X-axis: Horizontal, with positive values increasing rightward\n"177            "- Y-axis: Vertical, with positive values increasing upward\n"178            "- Z-axis: Depth, with positive values increasing toward the viewer\n"179            "- The origin point (0,0,0) is established at the centroid of the 3D mesh geometry.\n\n"180            "When analyzing rotations, consider the viewpoint from the positive end of the rotation axis."181        ),182        183        # Version 5 (More casual explanation)184        (185            "Let's define our 3D space like this:\n"186            "- The x-axis goes left to right (right is positive)\n"187            "- The y-axis goes bottom to top (up is positive)\n"188            "- The z-axis comes from the screen toward you (out is positive)\n"189            "- The zero point (0,0,0) sits exactly at the center of mass of the 3D object.\n\n"190            "When thinking about rotations, imagine you're looking along the axis from its positive end."191        ),192        193        # Version 6 (Formal academic style)194        (195            "The spatial reference frame utilized herein consists of:\n"196            "- An x-axis oriented horizontally (positive rightward)\n"197            "- A y-axis oriented vertically (positive upward)\n"198            "- A z-axis oriented perpendicular to the viewing plane (positive outward)\n"199            "- An origin (0,0,0) that is coincident with the geometric centroid of the object mesh.\n\n"200            "Rotational transformations are visualized from the perspective of an observer positioned along the positive direction of the axis in question."201        )202    ]203    coordinate_system = random.choice(coordinate_system_templates)204    angle_constraints_templates = [205        # Version 1206        (207            f"The angle of rotation must be one of these specific values: {possible_angles} degrees. "208            f"Positive angles denote clockwise rotation when viewed along the axis's positive direction."209        ),210 211        # Version 2212        (213            f"Permitted rotation angles are limited to these values: {possible_angles}. "214            f"The convention used is: a positive angle signifies rotation in the clockwise direction, assuming a viewpoint looking along the positive axis."215        ),216 217        # Version 3218        (219            f"Rotation is constrained to these angles: {possible_angles} degrees. "220            f"A positive angle value corresponds to a clockwise turn relative to an observer looking along the positive direction of the rotation axis."221        ),222    ]223    angle_constraints = random.choice(angle_constraints_templates)224    225    # Add examples text if examples are provided226    example_text = ""227    if examples and len(examples) > 0:228        example_text = "\n### EXAMPLES OF ROTATION ###\n"229        for idx, (_, example_angle) in enumerate(examples):230            if generation_mode == "combined":231                img_num = idx + 1232                example_text += f"\nExample {idx+1}: Image {img_num} shows a 3D object with its left half showing the initial view and right half showing a {example_angle} degree rotation around the {axis}-axis.\n"233            else:  # separate mode234                img_start = idx * 2 + 1235                img_end = idx * 2 + 2236                example_text += f"\nExample {idx+1}: Image {img_start} shows the initial view and Image {img_end} shows the object after a {example_angle} degree rotation around the {axis}-axis.\n"237    238    # Different instructions based on difficulty239    if difficulty == "easy":240        # For easy mode - axis is provided, internal reasoning but only output number241        thinking_instructions_templates = [242            # Version 1: More concise, action-oriented243            (244                f"CRITICAL STEPS for finding the rotation angle:"245                f"\n\n1. Analyze Object Structure: Examine both views thoroughly to understand the object's form."246                f"\n\n2. Evaluate ALL {axis}-Axis Options: You must consider every angle in this list: {possible_angles}."247                f"\n   - Visualize Rotation: For each angle, mentally rotate the object around the {axis}-axis."248                f"\n   - Compare to Target View: Check how each visualized rotation matches the second view."249                f"\n   - Complete Evaluation First: Do not choose an angle until all in {possible_angles} are tested."250                f"\n\n3. Select Best Match: After reviewing all possibilities, pick the angle that correctly transforms the first view to the second."251                f"\n\n4. Final Verification: Mentally apply your chosen rotation one last time to confirm it perfectly matches the second view."252            ),253 254            # Version 2: Slightly more explanatory, guiding tone255            (256                f"Follow this methodical process to identify the correct rotation:"257                f"\n\n1. Understand the Object: Start by carefully studying the object's features in both views to grasp its 3D shape."258                f"\n\n2. Systematic Angle Check ({axis}-axis): It's essential to evaluate the full set of potential rotation angles: {possible_angles}."259                f"\n   - For every angle listed: Imagine rotating the object around the {axis}-axis by that specific amount."260                f"\n   - Match Visualization to Reality: Compare your mental image after rotation with the provided second view."261                f"\n   - Avoid Premature Decisions: Ensure you have mentally tested *all* angles before making a selection."262                f"\n\n3. Determine the Correct Angle: Once all angles ({possible_angles}) have been considered, choose the one rotation that best explains the change between views."263                f"\n\n4. Confirm Your Answer: As a final check, mentally perform the chosen rotation again to ensure it accurately produces the second view."264            ),265 266            # Version 3: Focus on comparison and elimination267            (268                f"Use this procedure to pinpoint the rotation angle accurately:"269                f"\n\n1. Initial Analysis: Compare the first and second views to understand the object's spatial configuration."270                f"\n\n2. Exhaustive {axis}-Axis Evaluation: You are required to assess each of these candidate rotation angles: {possible_angles}."271                f"\n   - Test Each Angle: Mentally simulate rotating the object around the {axis}-axis by each angle in the list."272                f"\n   - Cross-Reference Views: Evaluate how closely each simulated rotation aligns with the actual second view."273                f"\n   - Full Assessment Required: Withhold judgment until every single angle from {possible_angles} has been assessed."274                f"\n\n3. Identify the Matching Rotation: After assessing all options, select the angle that precisely transforms the first view into the second."275                f"\n\n4. Validate Your Choice: Double-check by mentally applying the selected rotation to confirm it yields the exact second view."276            )277        ]278        thinking_instructions = random.choice(thinking_instructions_templates)279        280        # Updated response format to match rot_pred_sft.py281        response_format = (282            f"IMPORTANT: You must ONLY output the rotation angle as a number from this list: {possible_angles}. "283            f"Your output should contain ONLY the number. "284            f"Do NOT include any reasoning, explanation, or additional text - ONLY the number."285            f"\n\nExample of correct output format: 30"286            f"\n\nIncorrect output formats:"287            f"\n\"I think it's 30 degrees\""288            f"\n\"The rotation angle is 30\""289            f"\n\"30 degrees\""290        )291 292        task_description = (293            f"Your task is to determine the angle of rotation around the {axis}-axis in degrees."294        )295        296    else:  # hard mode - axis is not provided297        thinking_instructions = (298            f"IMPORTANT: Please follow this systematic approach to determine the rotation:"299            f"\n\n1. First, analyze the object's features in both views to understand its structure."300            f"\n\n2. Consider what would happen if rotation occurred around each of the three axes (x, y, and z):"301            f"\n   - For x-axis rotation: What specific features would change and how?"302            f"\n   - For y-axis rotation: What specific features would change and how?"303            f"\n   - For z-axis rotation: What specific features would change and how?"304            f"\n   - Based on the observed changes, explain which axis makes the most sense and why."305            f"\n\n3. Once you've determined the most likely axis, evaluate ALL of these possible rotation angles: {possible_angles}"306            f"\n   - For each angle in the list, describe what the object would look like after rotating around your chosen axis by that amount"307            f"\n   - Compare these descriptions with the actual second view"308            f"\n   - DO NOT make a decision until you have evaluated all angles in the list"309            f"\n\n4. After evaluating all angles, choose the one that best matches the observed changes"310        )311        312        response_format = (313            f"Place your detailed reasoning process in <thinking></thinking> tags. Your reasoning should include:"314            f"\n- Analysis of how rotation around each axis would affect the object"315            f"\n- Systematic evaluation of possible rotation angles from the provided list"316            f"\n- Specific visual features you used to determine your answer"317            f"\n\nThen provide your final answer in <rotation_axis></rotation_axis> and <rotation_angle></rotation_angle> tags respectively (use only x, y, or z for axis and only a number from the list for angle)."318            f"\ni.e., <thinking> your reasoning process here </thinking><rotation_axis> your predicted axis here </rotation_axis><rotation_angle> your predicted degrees here </rotation_angle>"319        )320        321        322        # task_description = (323        #     f"Your task is to determine which axis the object was rotated around and by what angle."324        # )325                326 327        task_description_templates = [328            "Identify the object's axis of rotation and the corresponding angle.",329            "You need to figure out around which axis the object turned, and by how much.",330            "Ascertain the rotational axis and the magnitude of the angle applied to the object.",331            "Find out both the specific axis used for the object's rotation and the degree of that rotation.",332            "The objective is to specify the rotation parameters for the object: its axis and angle."333        ]334        task_description = random.choice(task_description_templates)335    # Generate the prompt based on generation mode336    if generation_mode == "combined":337        test_img_num = len(examples) + 1 if examples else 1338        prompt = (339            f"IMPORTANT: I'm showing you {len(examples) + 1 if examples else 1} image{'s' if examples else ''} of 3D objects. "340            f"{'Each' if examples else 'The'} image contains TWO separate 3D renderings side-by-side. "  # Changed example to examples341            f"\n\nThe LEFT HALF shows a 3D object in its initial orientation. "342            f"The RIGHT HALF shows the SAME 3D object after being rotated."343            f"\n\n{task_description}"344            f"\n\n{coordinate_system}"345            f"\n\n{angle_constraints}"346            f"\n\n{example_text}"347            f"\n\n### YOUR TASK ###"348            f"\nNow, for Image {test_img_num}, determine the angle of rotation around the {axis}-axis."349            f"\n{'' if not examples else 'Based on the example provided, '}analyze Image {test_img_num} carefully."  # Changed example to examples350            f"\n\n{thinking_instructions}"351            f"\n\n{response_format}"352        )353    elif generation_mode == "separate_shuffle":354        test_img_start = len(examples) * 2 + 1 if examples else 1355        test_img_end = len(examples) * 2 + 2 if examples else 2356        begin_description = (357            f"I'm showing you {len(examples) * 2 + 2 if examples else 2} images of 3D objects. "358            f"{'For each example or test case, ' if examples else ''}two images represent the same object before and after rotation."  # Changed example to examples359        )360        end_description = (361            f"\n\n### YOUR TASK ###"362            f"\nNow, determine the angle of rotation around the {axis}-axis from Image {test_img_start} to Image {test_img_end}."363            f"\n{'' if not examples else 'Based on the example provided, '}analyze the rotation carefully."  # Changed example to examples364            f"\n\n{thinking_instructions}"365            f"\n\n{response_format}"366        )367        368        prompt_list = [task_description, coordinate_system, angle_constraints, example_text]369        random.shuffle(prompt_list)370        prompt = begin_description + "\n\n".join(prompt_list) + end_description371    elif generation_mode == "separate":372        # Calculate image numbers based on examples373        test_img_start = len(examples) * 2 + 1 if examples else 1374        test_img_end = len(examples) * 2 + 2 if examples else 2375        prompt = (376            f"I'm showing you {len(examples) * 2 + 2 if examples else 2} images of 3D objects. "377            f"{'For each example or test case, ' if examples else ''}two images represent the same object before and after rotation."  # Changed example to examples378            f"\n\n{task_description}"379            f"\n\n{coordinate_system}"380            f"\n\n{angle_constraints}"381            f"\n\n{example_text}"382            f"\n\n### YOUR TASK ###"383            f"\nNow, determine the angle of rotation around the {axis}-axis from Image {test_img_start} to Image {test_img_end}."384            f"\n{'' if not examples else 'Based on the example provided, '}analyze the rotation carefully."  # Changed example to examples385            f"\n\n{thinking_instructions}"386            f"\n\n{response_format}"387        )388    else:389        raise ValueError(f"Invalid generation mode: {generation_mode}")390    return prompt391 392def create_metadata_jsonl_separate(input_dir, output_file, example_dir=None, possible_angles=[45, 315], difficulty="easy", max_examples=2):393    """Create metadata JSONL file for all images in input_dir (combined mode)"""394    # Get all PNG files in the input directory395    png_files = glob.glob(os.path.join(input_dir, "*.png"))396    397    # Sort files to ensure consistent order398    png_files = sorted(png_files)399    400    if not png_files:401        print(f"No PNG files found in {input_dir}")402        return403    404    print(f"Found {len(png_files)} PNG files in {input_dir}")405    406    # Load and organize examples if example_dir is provided407    organized_examples = None408    if example_dir:409        examples = load_examples(example_dir, "combined")410        organized_examples = organize_examples(examples, "combined")411    412    # Create output directory if it doesn't exist413    output_dir = os.path.dirname(output_file)414    os.makedirs(output_dir, exist_ok=True)415    416    # Process each file and create metadata entries417    entries = []418    419    for png_file in tqdm(png_files, desc="Creating metadata for combined mode"):420        # Parse ground truth from filename421        axis, angle = parse_ground_truth(os.path.basename(png_file))422        423        if axis is None or angle is None:424            print(f"Skipping {png_file} - could not parse ground truth")425            continue426        427        # Get the relative path to the image428        rel_path = os.path.relpath(png_file, os.path.dirname(output_file))429        430        # Generate a unique ID based on the filename431        image_base_id = os.path.splitext(os.path.basename(png_file))[0]432        433        # Select an example if examples are available434        examples = None435        if organized_examples:436            examples = select_examples(organized_examples, axis, possible_angles, max_examples)437        438        # Construct prompt with examples439        prompt = construct_prompt_with_examples(axis, possible_angles, examples, difficulty, generation_mode="combined")440 441        # Create assistant response based on difficulty442        if difficulty == "easy":443            # For easy mode, just output the number444            assistant_content = f"{angle}"445        else:446            # For hard mode, include both axis and angle in XML tags447            assistant_content = f"<thinking>Detailed reasoning about rotation axis and angle...</thinking><rotation_axis>{axis}</rotation_axis><rotation_angle>{angle}</rotation_angle>"448        449        # Create the conversations array450        conversations = []451        452        # Fix the human message construction to handle multiple examples453        human_value = ""454        455        # Add example images if available456        if examples:457            for example_path, _ in examples:458                example_rel_path = os.path.relpath(example_path, os.path.dirname(output_file))459                human_value += f"<image>{example_rel_path}</image>\n"460        461        # Add test image462        human_value += f"<image>{rel_path}</image>\n{prompt}"463        464        conversations.append({465            "from": "human",466            "value": human_value467        })468        469        # Add assistant response470        conversations.append({471            "from": "gpt",472            "value": assistant_content473        })474        475        # Create entry with the correct format476        entry = {477            "id": image_base_id,478            "image": rel_path,479            "conversations": conversations480        }481        482        entries.append(entry)483    484    # Write entries to JSONL file485    with open(output_file, 'w') as f:486        for entry in entries:487            f.write(json.dumps(entry) + '\n')488    489    print(f"\nSummary for combined mode:")490    print(f"  Found {len(png_files)} PNG files")491    print(f"  Created metadata for {len(entries)} entries")492    print(f"  Output file: {output_file}")493 494def create_metadata_jsonl_separate(input_dir, output_file, example_dir=None, possible_angles=[45, 315], difficulty="easy", max_examples=2, generation_mode="separate"):495    """Create metadata JSONL file for folders in input_dir (separate mode)"""496    # Get all directories in the input directory497    folders = [f for f in glob.glob(os.path.join(input_dir, "*")) 498              if os.path.isdir(f) and os.path.basename(f) != "examples"]499    500    # Sort folders to ensure consistent order501    folders = sorted(folders)502    503    if not folders:504        print(f"No folders found in {input_dir}")505        return506    507    print(f"Found {len(folders)} folders in {input_dir}")508    509    # Load and organize examples if example_dir is provided510    organized_examples = None511    if example_dir:512        examples = load_examples(example_dir, "separate")513        organized_examples = organize_examples(examples, "separate")514    515    # Create output directory if it doesn't exist516    output_dir = os.path.dirname(output_file)517    os.makedirs(output_dir, exist_ok=True)518    519    # Process each folder and create metadata entries520    entries = []521    valid_folders = 0522    523    for folder in tqdm(folders, desc="Creating metadata for separate mode"):524        folder_name = os.path.basename(folder)525        526        # Parse ground truth from folder name527        axis, angle = parse_ground_truth(folder_name)528        529        if axis is None or angle is None:530            print(f"Skipping {folder} - could not parse ground truth")531            continue532        533        # Check for the two required images in the folder534        ini_path = os.path.join(folder, f"{folder_name}_ini.png")535        rot_path = os.path.join(folder, f"{folder_name}_rot.png")536        537        if not os.path.exists(ini_path):538            print(f"Skipping {folder} - missing initial view image")539            continue540            541        if not os.path.exists(rot_path):542            print(f"Skipping {folder} - missing rotated view image")543            continue544        545        # Get the relative paths to the images546        rel_ini_path = os.path.relpath(ini_path, os.path.dirname(output_file))547        rel_rot_path = os.path.relpath(rot_path, os.path.dirname(output_file))548        549        # Select an example if examples are available550        examples = None551        if organized_examples:552            examples = select_examples(organized_examples, axis, possible_angles, max_examples)553        554        # Update this to construct_prompt_with_examples555        prompt = construct_prompt_with_examples(axis, possible_angles, examples, difficulty, generation_mode=generation_mode)556        557        # Create assistant response based on difficulty558        if difficulty == "easy":559            # For easy mode, just output the number560            assistant_content = f"{angle}"561        else:562            # For hard mode, include both axis and angle in XML tags563            assistant_content = f"<thinking>Detailed reasoning about rotation axis and angle...</thinking><rotation_axis>{axis}</rotation_axis><rotation_angle>{angle}</rotation_angle>"564        565        # Create the conversations array566        conversations = []567        568        # Prepare images array for the entry569        all_image_paths = []570 571        # Add example images if available572        if examples:573            for example_folder, _ in examples:574                example_folder_name = os.path.basename(example_folder)575                example_ini_path = os.path.join(example_folder, f"{example_folder_name}_ini.png")576                example_rot_path = os.path.join(example_folder, f"{example_folder_name}_rot.png")577                578                example_rel_ini_path = os.path.relpath(example_ini_path, os.path.dirname(output_file))579                example_rel_rot_path = os.path.relpath(example_rot_path, os.path.dirname(output_file))580                581                all_image_paths.append(example_rel_ini_path)582                all_image_paths.append(example_rel_rot_path)583        584        # Add test images585        all_image_paths.append(rel_ini_path)586        all_image_paths.append(rel_rot_path)587        588        # Update the human message tags to match number of images589        # For 2 examples (4 images) + 2 test images = 6 total images590        image_tags = "<image>\n" * len(all_image_paths)591        human_value = image_tags + prompt592        593        conversations.append({594            "from": "human",595            "value": human_value596        })597        598        # Add assistant response599        conversations.append({600            "from": "gpt",601            "value": assistant_content602        })603        604        # Create entry with the correct format605        entry = {606            "id": folder_name,607            "image": all_image_paths,608            "conversations": conversations609        }610        611        entries.append(entry)612        valid_folders += 1613    614    # Write entries to JSONL file615    with open(output_file, 'w') as f:616        for entry in entries:617            f.write(json.dumps(entry) + '\n')618    619    print(f"\nSummary for separate mode:")620    print(f"  Found {len(folders)} folders")621    print(f"  Created metadata for {valid_folders} valid folders")622    print(f"  Output file: {output_file}")623    624def main():625    parser = argparse.ArgumentParser(description="Create metadata JSONL for rotation dataset")626    parser.add_argument('--input-dir', type=str, required=True,627                      help="Directory containing rotation dataset images or folders")628    parser.add_argument('--output-file', type=str, default="metadata.jsonl",629                      help="Output JSONL file path")630    parser.add_argument('--example-dir', type=str, default=None,631                      help="Directory containing example images for in-context learning")632    parser.add_argument('--possible-angles', type=int, nargs='+', default=[45, 315],633                    help="List of possible rotation angles in degrees (e.g., 45 315)")634    parser.add_argument('--difficulty', type=str, choices=["easy", "hard"], default="easy",635                      help="Difficulty mode: easy (axis provided) or hard (axis not provided)")636    parser.add_argument('--generation-mode', type=str, choices=["combined", "separate", "separate_shuffle"], default="combined",637                      help="Mode for dataset generation (combined = one image with both views, separate = folder with two images, separate_shuffle = separate with shuffled prompt sections)")638    parser.add_argument('--random-seed', type=int, default=None,639                      help="Random seed for example selection (None for true randomness)")640    parser.add_argument('--max-examples', type=int, default=1,641                      help="Maximum number of examples to include for each test case (default: 1)")642    643    args = parser.parse_args()644    645    # Set random seed for reproducibility if provided646    if args.random_seed is not None:647        print(f"Using fixed random seed: {args.random_seed}")648        random.seed(args.random_seed)649    else:650        print("Using true randomness (different examples each run)")651    652    print(f"Creating metadata JSONL for rotation dataset:")653    print(f"Input directory: {args.input_dir}")654    print(f"Output file: {args.output_file}")655    656    if args.example_dir:657        print(f"Example directory: {args.example_dir}")658    659    print(f"Possible angles: {args.possible_angles}")660    print(f"Difficulty mode: {args.difficulty}")661    print(f"Generation mode: {args.generation_mode}")662    663    # Check if example_dir is None but there's an 'examples' subdirectory in input_dir664    if args.example_dir is None and os.path.exists(os.path.join(args.input_dir, "examples")):665        args.example_dir = os.path.join(args.input_dir, "examples")666        print(f"Using examples directory: {args.example_dir}")667    668    if args.generation_mode == "combined":669        create_metadata_jsonl_combined(670            input_dir=args.input_dir,671            output_file=args.output_file,672            example_dir=args.example_dir,673            possible_angles=args.possible_angles,674            difficulty=args.difficulty,675            max_examples=args.max_examples  # Make sure this is passed properly676        )677    else:  # separate or separate_shuffle mode678        create_metadata_jsonl_separate(679            input_dir=args.input_dir,680            output_file=args.output_file,681            example_dir=args.example_dir,682            possible_angles=args.possible_angles,683            difficulty=args.difficulty,684            max_examples=args.max_examples,685            generation_mode=args.generation_mode  # Pass the actual mode686        )687 688if __name__ == "__main__":689    main()