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()