SCMayS/hydata

1import os2import re3import glob4import json5import argparse6from tqdm import tqdm7from pathlib import Path8 9def parse_ground_truth(name):10    """Extract ground truth rotation axis and angle from filename or folder name"""11    # Remove file extension if present12    basename = name.split(".")[0] if "." in name else name13    14    parts = basename.split("_")15    if len(parts) >= 4:  # figXXXX_XXX_axis_angle16        rotation_axis = parts[-2]  # Second to last element is axis17        rotation_angle = int(parts[-1])  # Last element is angle18        19        # Convert negative angles to 0-360 range20        if rotation_angle < 0:21            rotation_angle += 36022            23        return rotation_axis, rotation_angle24    25    print(f"Warning: Could not parse name: {basename}")26    return None, None27 28def construct_prompt(axis, angle_increment, difficulty="easy", generation_mode="combined"):29    """Create prompt for the VLM based on difficulty mode and generation mode"""30    # Generate list of all possible rotation angles based on angle increment31    possible_angles = []32    current_angle = 0 + angle_increment33    while current_angle < 360:34        possible_angles.append(current_angle)35        current_angle += angle_increment36    37    # Common instructions for both modes38    coordinate_system = (39        f"The 3D Cartesian coordinate system is defined as follows: "40        f"\n- x-axis: points horizontally from left to right (positive direction is right)"41        f"\n- y-axis: points vertically from bottom to top (positive direction is up)"42        f"\n- z-axis: points from inside the image toward the viewer (positive direction is out of the screen)"43        f"\n\nWhen discussing rotations around an axis, imagine looking along the positive direction of that axis (as if looking from the origin toward the positive end)."44    )45    46    angle_constraints = (47        f"The rotation angle is always a multiple of {angle_increment} degrees between 0 and 360 degrees inclusive. "48        f"A positive angle means rotation in the CLOCKWISE direction when looking along the positive direction of the axis. "49    )50    51    # Different instructions based on difficulty52    if difficulty == "easy":53        # For easy mode - axis is provided54        thinking_instructions = (55            f"IMPORTANT: Please follow this systematic approach to determine the rotation angle:"56            f"\n\n1. First, analyze the object's features in both views to understand its structure."57            f"\n\n2. For the {axis}-axis rotation, you must evaluate ALL of these possible rotation angles: {possible_angles}"58            f"\n   - For each angle in the list, describe what the object would look like after rotating around the {axis}-axis by that amount"59            f"\n   - Compare these descriptions with the actual second view"60            f"\n   - DO NOT make a decision until you have evaluated all possible angles in the list"61            f"\n\n3. After evaluating all angles, choose the one that best matches the observed changes"62            f"\n\n4. Verify your answer by mentally applying the rotation to confirm it matches the second view"63        )64        65        response_format = (66            f"Place your detailed reasoning process in <think></think> tags. Your reasoning should include:"67            f"\n- Systematic evaluation of possible rotation angles from the provided list"68            f"\n- Specific visual features you used to determine your answer"69            f"\n\nThen provide your final answer in <rotation_angle></rotation_angle> tags (use only a number from the list for angle)."70            f"\ni.e., <think> your reasoning process here </think><rotation_angle> your predicted degrees here </rotation_angle>"71        )72        73        task_description = (74            f"Your task is to determine the angle of rotation around the {axis}-axis in degrees."75        )76        77    else:  # hard mode - axis is not provided78        thinking_instructions = (79            f"IMPORTANT: Please follow this systematic approach to determine the rotation:"80            f"\n\n1. First, analyze the object's features in both views to understand its structure."81            f"\n\n2. Consider what would happen if rotation occurred around each of the three axes (x, y, and z):"82            f"\n   - For x-axis rotation: What specific features would change and how?"83            f"\n   - For y-axis rotation: What specific features would change and how?"84            f"\n   - For z-axis rotation: What specific features would change and how?"85            f"\n   - Based on the observed changes, explain which axis makes the most sense and why."86            f"\n\n3. Once you've determined the most likely axis, evaluate ALL of these possible rotation angles: {possible_angles}"87            f"\n   - For each angle in the list, describe what the object would look like after rotating around your chosen axis by that amount"88            f"\n   - Compare these descriptions with the actual second view"89            f"\n   - DO NOT make a decision until you have evaluated all angles in the list"90            f"\n\n4. After evaluating all angles, choose the one that best matches the observed changes"91        )92        93        response_format = (94            f"Place your detailed reasoning process in <thinking></thinking> tags. Your reasoning should include:"95            f"\n- Analysis of how rotation around each axis would affect the object"96            f"\n- Systematic evaluation of possible rotation angles from the provided list"97            f"\n- Specific visual features you used to determine your answer"98            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)."99            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>"100        )101        102        task_description = (103            f"Your task is to determine which axis the object was rotated around and by what angle."104        )105    106    # Generate the prompt based on generation mode107    if generation_mode == "combined":108        prompt = (109            f"IMPORTANT: The image I'm showing you contains TWO separate 3D renderings side-by-side. "110            f"This is a single image file with two distinct rendered views placed next to each other. "111            f"\n\nThe LEFT HALF shows a 3D object in its initial orientation. "112            f"The RIGHT HALF shows the SAME 3D object after being rotated. "113            f"\n\n{task_description}"114            f"\n\n{coordinate_system}"115            f"\n\n{angle_constraints}"116            f"\n\n{thinking_instructions}"117            f"\n\n{response_format}"118        )119    else:  # separate mode120        prompt = (121            f"I'm showing you TWO images of the same 3D object. "122            f"The FIRST image shows the object in its initial orientation. "123            f"The SECOND image shows the object after being rotated. "124            f"\n\n{task_description}"125            f"\n\n{coordinate_system}"126            f"\n\n{angle_constraints}"127            f"\n\n{thinking_instructions}"128            f"\n\n{response_format}"129        )130            131    return prompt132 133def create_metadata_jsonl_combined(input_dir, output_file, angle_increment=30, difficulty="easy"):134    """Create metadata JSONL file for all images in input_dir (combined mode)"""135    # Get all PNG files in the input directory136    png_files = glob.glob(os.path.join(input_dir, "*.png"))137    138    # Sort files to ensure consistent order139    png_files = sorted(png_files)140    141    if not png_files:142        print(f"No PNG files found in {input_dir}")143        return144    145    print(f"Found {len(png_files)} PNG files in {input_dir}")146    147    # Create output directory if it doesn't exist148    output_dir = os.path.dirname(output_file)149    os.makedirs(output_dir, exist_ok=True)150    151    # Get the last folder name from the input directory152    last_folder = os.path.basename(os.path.normpath(input_dir))153    154    # Define the target base directory155    target_base_dir = f"/lustre/fsw/portfolios/av/users/shiyil/yunfei/MM-EUREKA/data/{last_folder}"156    157    # Process each file and create metadata entries158    entries = []159    160    for png_file in tqdm(png_files, desc="Creating metadata for combined mode"):161        # Parse ground truth from filename162        axis, angle = parse_ground_truth(os.path.basename(png_file))163        164        if axis is None or angle is None:165            print(f"Skipping {png_file} - could not parse ground truth")166            continue167        168        # Get the basename169        basename = os.path.basename(png_file)170        171        # Create the new image path for the target system172        target_image_path = f"{target_base_dir}/{basename}"173        174        # Construct prompt for combined mode175        prompt = construct_prompt(axis, angle_increment, difficulty, generation_mode="combined")176        177        # Create answer format based on difficulty WITH image path178        if difficulty == "easy":179            # For easy mode, only include angle in the answer (axis is provided in the prompt)180            answer = f"<angle>{angle}</angle><image_path>{target_image_path}</image_path>"181        else:182            # For hard mode, include both axis and angle183            answer = f"<axis>{axis}</axis><angle>{angle}</angle><image_path>{target_image_path}</image_path>"184        185        # Create entry with the modified image path186        entry = {187            "message": json.dumps([{188                "role": "user",189                "content": [190                    {191                        "type": "image",192                        "image": target_image_path193                    },194                    {195                        "type": "text",196                        "text": prompt197                    }198                ]199            }]),200            "answer": answer201        }202        203        entries.append(entry)204    205    # Write entries to JSONL file206    with open(output_file, 'w') as f:207        for entry in entries:208            f.write(json.dumps(entry) + '\n')209    210    print(f"\nSummary for combined mode:")211    print(f"  Found {len(png_files)} PNG files")212    print(f"  Created metadata for {len(entries)} entries")213    print(f"  Output file: {output_file}")214    print(f"  Image paths are set to: {target_base_dir}/[filename].png")215 216def create_metadata_jsonl_separate(input_dir, output_file, angle_increment=30, difficulty="easy"):217    """Create metadata JSONL file for folders in input_dir (separate mode)"""218    # Get all directories in the input directory219    folders = [f for f in glob.glob(os.path.join(input_dir, "*")) if os.path.isdir(f)]220    221    # Sort folders to ensure consistent order222    folders = sorted(folders)223    224    if not folders:225        print(f"No folders found in {input_dir}")226        return227    228    print(f"Found {len(folders)} folders in {input_dir}")229    230    # Create output directory if it doesn't exist231    output_dir = os.path.dirname(output_file)232    os.makedirs(output_dir, exist_ok=True)233    234    # Get the last folder name from the input directory235    last_folder = os.path.basename(os.path.normpath(input_dir))236    237    # Define the target base directory238    target_base_dir = f"/lustre/fsw/portfolios/av/users/shiyil/yunfei/MM-EUREKA/data/{last_folder}"239    240    # Process each folder and create metadata entries241    entries = []242    valid_folders = 0243    244    for folder in tqdm(folders, desc="Creating metadata for separate mode"):245        folder_name = os.path.basename(folder)246        247        # Parse ground truth from folder name248        axis, angle = parse_ground_truth(folder_name)249        250        if axis is None or angle is None:251            print(f"Skipping {folder} - could not parse ground truth")252            continue253        254        # Check for the two required images in the folder255        ini_path = os.path.join(folder, f"{folder_name}_ini.png")256        rot_path = os.path.join(folder, f"{folder_name}_rot.png")257        258        if not os.path.exists(ini_path):259            print(f"Skipping {folder} - missing initial view image")260            continue261            262        if not os.path.exists(rot_path):263            print(f"Skipping {folder} - missing rotated view image")264            continue265        266        # Create target paths for remote system267        target_folder_path = f"{target_base_dir}/{folder_name}"268        target_ini_path = f"{target_base_dir}/{folder_name}/{folder_name}_ini.png"269        target_rot_path = f"{target_base_dir}/{folder_name}/{folder_name}_rot.png"270        271        # Construct prompt for separate mode272        prompt = construct_prompt(axis, angle_increment, difficulty, generation_mode="separate")273        274        # Create answer format based on difficulty WITH folder path275        if difficulty == "easy":276            # For easy mode, only include angle in the answer (axis is provided in the prompt)277            answer = f"<angle>{angle}</angle><image_path>{target_folder_path}</image_path>"278        else:279            # For hard mode, include both axis and angle280            answer = f"<axis>{axis}</axis><angle>{angle}</angle><image_path>{target_folder_path}</image_path>"281        282        # Create entry with both image paths283        entry = {284            "message": json.dumps([{285                "role": "user",286                "content": [287                    {"type": "image", "image": target_ini_path},288                    {"type": "image", "image": target_rot_path},289                    {"type": "text", "text": prompt}290                ]291            }]),292            "answer": answer293        }294        295        entries.append(entry)296        valid_folders += 1297    298    # Write entries to JSONL file299    with open(output_file, 'w') as f:300        for entry in entries:301            f.write(json.dumps(entry) + '\n')302    303    print(f"\nSummary for separate mode:")304    print(f"  Found {len(folders)} folders")305    print(f"  Created metadata for {valid_folders} valid folders")306    print(f"  Output file: {output_file}")307    print(f"  Image paths format: {target_base_dir}/[folder_name]/[folder_name]_[ini/rot].png")308 309def main():310    parser = argparse.ArgumentParser(description="Create metadata JSONL for rotation dataset")311    parser.add_argument('--input-dir', type=str, required=True,312                      help="Directory containing rotation dataset images or folders")313    parser.add_argument('--output-file', type=str, default="rotation_metadata.jsonl",314                      help="Output JSONL file path")315    parser.add_argument('--angle-increment', type=int, default=30,316                      help="Angle increment used in the dataset (e.g., 30, 45, 90)")317    parser.add_argument('--difficulty', type=str, choices=["easy", "hard"], default="easy",318                      help="Difficulty mode: easy (axis provided) or hard (axis not provided)")319    parser.add_argument('--generation-mode', type=str, choices=["combined", "separate"], default="combined",320                      help="Mode for dataset generation (combined = one image with both views, separate = folder with two images)")321    322    args = parser.parse_args()323    324    print(f"Creating metadata JSONL for rotation dataset:")325    print(f"Input directory: {args.input_dir}")326    print(f"Output file: {args.output_file}")327    print(f"Angle increment: {args.angle_increment} degrees")328    print(f"Difficulty mode: {args.difficulty}")329    print(f"Generation mode: {args.generation_mode}")330    331    if args.generation_mode == "combined":332        create_metadata_jsonl_combined(333            input_dir=args.input_dir,334            output_file=args.output_file,335            angle_increment=args.angle_increment,336            difficulty=args.difficulty337        )338    else:  # separate mode339        create_metadata_jsonl_separate(340            input_dir=args.input_dir,341            output_file=args.output_file,342            angle_increment=args.angle_increment,343            difficulty=args.difficulty344        )345 346if __name__ == "__main__":347    main()