Spaces:
Running
Running
File size: 3,680 Bytes
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 |
import os
import re
import shutil
from config import MODEL_OUTPUT_IMAGE_NAMES
ROOT_DIR = "benchmark_images_generations"
OUTPUT_DIR = "data/"
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
# Outputs
if any([not os.path.exists(os.path.join(sample_dir_path, image_name)) for image_name in MODEL_OUTPUT_IMAGE_NAMES.values()]):
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 _, image_name in MODEL_OUTPUT_IMAGE_NAMES.items():
image_path = os.path.join(sample_dir_path, image_name)
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() |