restructure_folder.py

import os
import re
import shutil

from PIL import Image, ImageDraw

from config import MODEL_OUTPUT_IMAGE_NAMES

ROOT_DIR = "benchmark_images_generations"
OUTPUT_DIR = "data/"

def create_image_with_bbox_border(
    bg_image_path: str,
    mask_image_path: str,
    output_image_path: str,
    border_color: tuple = (245, 40, 145, 205),
    border_width: int = 5,
) -> None:
    """
    Opens a background image and a mask image (defining a rectangular bounding box),
    ensures they have the same size, then creates and saves an image showing the
    background with a semi-transparent border of the bounding box to the specified
    output path. No console output is produced by this function, and no image is displayed.

    Args:
        bg_image_path (str): Path to the background image.
        mask_image_path (str): Path to the mask image (expected to be black
                               with a white rectangle defining the box).
        output_image_path (str): Path where the resulting image will be saved.
        border_color (tuple): RGBA tuple for the border color.
                              The alpha component (0-255) controls transparency.
        border_width (int): Width of the bounding box border in pixels.

    Raises:
        FileNotFoundError: If either the background or mask image file is not found.
        ValueError: If no bounding box (non-black area) is found in the mask image.
        Exception: For other PIL-related errors during image processing (e.g.,
                   unsupported file format, issues during drawing, or saving errors).
                   This includes TypeError if the installed Pillow version does not
                   support the 'width' parameter for `ImageDraw.rectangle`.
    """
    # 1. Open background and mask images
    # These will raise FileNotFoundError if paths are invalid.
    bg_image = Image.open(bg_image_path)
    mask_image = Image.open(mask_image_path)

    # 2. Ensure images have the same size by resizing the mask to match the background
    if bg_image.size != mask_image.size:
        # Use Image.NEAREST to preserve sharp edges for the mask when resizing.
        mask_image = mask_image.resize(bg_image.size, Image.NEAREST)

    # 3. Convert background image to RGBA to allow drawing with transparency
    if bg_image.mode != 'RGBA':
        bg_image = bg_image.convert('RGBA')

    # 4. Find the bounding box from the mask image
    # Convert mask to grayscale ('L') to reliably get the bbox of the white area
    # (white pixels will be 255, black will be 0).
    mask_gray = mask_image.convert('L')
    
    # Get the bounding box of the non-black (i.e., white) region in the mask.
    # If the mask is all black, getbbox() returns None.
    bbox = mask_gray.getbbox()

    if bbox is None:
        raise ValueError("No bounding box found in the mask image (it might be entirely black or invalid).")

    # 5. Create a drawing context on the background image (which is now RGBA)
    draw = ImageDraw.Draw(bg_image)

    # 6. Draw the semi-transparent rectangle border
    # The 'width' parameter for draw.rectangle specifies the border thickness.
    # This requires a reasonably modern version of Pillow (5.2.0+).
    # If not supported, a TypeError will be raised by Pillow.
    draw.rectangle(bbox, outline=border_color, width=border_width)
        
    # 7. Save the resulting image to the specified output path
    # This can raise various exceptions (e.g., IOError, KeyError for unknown format)
    bg_image.save(output_image_path)




def main():
    for domain in os.listdir(ROOT_DIR):
        domain_dir = os.path.join(ROOT_DIR, domain)
        for i, sample_dir in enumerate(os.listdir(domain_dir)):
            if sample_dir == ".DS_Store":
                continue
            sample_dir_path = os.path.join(domain_dir, sample_dir)
            prompt = sample_dir[4:].strip()

            output_sample_dir = os.path.join(OUTPUT_DIR, domain, f"sample_{i}")
            os.makedirs(output_sample_dir, exist_ok=True)

            # --- Look for all the images needed to complete the sample ---
            
            # Input background
            ## look for an image (either .jpg or .png) whose name is composed as "bg{number}"
            ## and copy it to the output directory, then rename it to input_bg.jpg
            input_bg = None
            for file in os.listdir(sample_dir_path):
                if re.match(r"bg\d+\.(jpg|png)", file):
                    input_bg = file
                    break
            if input_bg:
                input_bg_path = os.path.join(sample_dir_path, input_bg)
                shutil.copy(input_bg_path, os.path.join(output_sample_dir, "input_bg.jpg"))
            else:
                print(f"Warning: No input background found in {sample_dir_path}. Skipping sample {i}...")
                continue

            # Input foreground
            ## look for an image (either .jpg or .png) whose name is either composed as "fg_{alphanumeric}.jpg" or "fg_{alphanumeric}.png",
            ## or "fg{number}_{alphanumeric}.jpg" or "fg{number}_{alphanumeric}.png", and never contains the word "mask"
            ## and copy it to the output directory, then rename it to input_fg.jpg
            input_fg = None
            for file in os.listdir(sample_dir_path):
                if re.match(r"fg(_\w+)?\.(jpg|png)", file) or re.match(r"fg\d+(_\w+)?\.(jpg|png)", file):
                    if "mask" not in file:
                        input_fg = file
                        break
            if input_fg:
                input_fg_path = os.path.join(sample_dir_path, input_fg)
                shutil.copy(input_fg_path, os.path.join(output_sample_dir, "input_fg.jpg"))
            else:
                print(f"Warning: No input foreground found in {sample_dir_path}. Skipping sample {i}...")
                continue

            # Input bb
            input_bb = None
            for file in os.listdir(sample_dir_path):
                if file == "mask_bg_fg.jpg" or file == "mask_bg_fg.png":
                    input_bb = file
                    break
            if input_bb and input_bg:
                # Create composed image
                create_image_with_bbox_border(
                    bg_image_path=os.path.join(output_sample_dir, "input_bg.jpg"),
                    mask_image_path=os.path.join(sample_dir_path, input_bb),
                    output_image_path=os.path.join(output_sample_dir, "input_bg_bb.png"),
                )

            # Outputs
            if not all(
                [
                    os.path.exists(os.path.join(sample_dir_path, "cp_bg_fg.jpg")),
                    os.path.exists(os.path.join(sample_dir_path, "kvedit.jpg")),
                    os.path.exists(os.path.join(sample_dir_path, "tf-icon.png")),
                ]
            ) or not any(
                [
                    # os.path.exists(os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0.png")),
                    # os.path.exists(os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0_vital-layers.png")),
                    os.path.exists(os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VFalse_taua0.4_taub0.8_guidance3.0_all-layers.png")),
                ]
            ):
                print(f"Warning: Not all output images found in {sample_dir_path}. Skipping sample {i}...")
                # Delete the output directory for this sample
                shutil.rmtree(output_sample_dir)
                continue

            for model_name, image_name in MODEL_OUTPUT_IMAGE_NAMES.items():
                # Patch needed because of inconsistency in the naming of the dit-editor images
                image_path = os.path.join(sample_dir_path, image_name)
                # if model_name != "dit-editor":
                #     image_path = os.path.join(sample_dir_path, image_name)
                # else:
                    # if os.path.exists(os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0.png")):
                    #     image_path = os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0.png")
                    # elif os.path.exists(os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0_vital-layers.png")):
                    #     image_path = os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0_vital-layers.png")
                    # else:
                    #     raise FileNotFoundError(f"Neither of the expected images for dit-editor found in {sample_dir_path}")

                target_path = os.path.join(output_sample_dir, image_name)
                if os.path.exists(image_path):
                    shutil.copy(image_path, target_path)
                else:
                    print(f"Warning: {image_name} not found in {sample_dir_path}. Skipping...")

            # -- Create the prompt.txt file --
            ## create a file named prompt.txt in the output directory and write the prompt in it
            prompt_file_path = os.path.join(output_sample_dir, "prompt.txt")
            with open(prompt_file_path, "w") as prompt_file:
                prompt_file.write(prompt)

if __name__ == "__main__":
    main()