Spaces:
Running
Running
File size: 9,468 Bytes
af5e0d4 2b6be77 af5e0d4 2b6be77 af5e0d4 2b6be77 af5e0d4 ed9aefd 2b6be77 ed9aefd af5e0d4 ed9aefd 2b6be77 ed9aefd af5e0d4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 |
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() |