File size: 8,622 Bytes
fe6327d |
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 |
import argparse
import shutil
from PIL import Image, ImageOps
import os
import numpy as np
from library.custom_logging import setup_logging
# Set up logging
log = setup_logging()
class ImageProcessor:
def __init__(self, input_folder, output_folder, group_size, include_subfolders, do_not_copy_other_files, pad, caption, caption_ext):
self.input_folder = input_folder
self.output_folder = output_folder
self.group_size = group_size
self.include_subfolders = include_subfolders
self.do_not_copy_other_files = do_not_copy_other_files
self.pad = pad
self.caption = caption
self.caption_ext = caption_ext
self.image_extensions = ('.png', '.jpg', '.jpeg', '.gif', '.webp', '.tiff')
def get_image_paths(self):
images = []
if self.include_subfolders:
for dirpath, dirnames, filenames in os.walk(self.input_folder):
for filename in filenames:
if filename.endswith(self.image_extensions):
images.append(os.path.join(dirpath, filename))
else:
images = [os.path.join(self.input_folder, f) for f in os.listdir(self.input_folder) if f.endswith(self.image_extensions)]
return images
def group_images(self, images):
sorted_images = sorted(images, key=lambda path: Image.open(path).size[0] / Image.open(path).size[1])
groups = [sorted_images[i:i+self.group_size] for i in range(0, len(sorted_images), self.group_size)]
return groups
def process_group(self, group, group_index):
if len(group) > 0:
aspect_ratios = self.get_aspect_ratios(group)
avg_aspect_ratio = np.mean(aspect_ratios)
if self.pad:
padded_images = self.pad_images(group, avg_aspect_ratio)
self.resize_and_save_images(padded_images, group_index, group)
else:
cropped_images = self.crop_images(group, avg_aspect_ratio)
self.resize_and_save_images(cropped_images, group_index, group)
if not self.do_not_copy_other_files:
self.copy_other_files(group, group_index)
def get_aspect_ratios(self, group):
aspect_ratios = []
for path in group:
with Image.open(path) as img:
width, height = img.size
aspect_ratios.append(width / height)
return aspect_ratios
def crop_images(self, group, avg_aspect_ratio):
cropped_images = []
for j, path in enumerate(group):
with Image.open(path) as img:
log.info(f" Processing image {j+1}: {path}")
img = self.crop_image(img, avg_aspect_ratio)
cropped_images.append(img)
return cropped_images
def crop_image(self, img, avg_aspect_ratio):
img_aspect_ratio = img.width / img.height
if img_aspect_ratio > avg_aspect_ratio:
# Too wide, reduce width
new_width = avg_aspect_ratio * img.height
left = (img.width - new_width) / 2
right = left + new_width
img = img.crop((left, 0, right, img.height))
else:
# Too tall, reduce height
new_height = img.width / avg_aspect_ratio
top = (img.height - new_height) / 2
bottom = top + new_height
img = img.crop((0, top, img.width, bottom))
return img
def resize_and_save_images(self, cropped_images, group_index, source_paths):
max_width = max(img.width for img in cropped_images)
max_height = max(img.height for img in cropped_images)
for j, img in enumerate(cropped_images):
img = img.resize((max_width, max_height))
os.makedirs(self.output_folder, exist_ok=True)
original_filename = os.path.basename(source_paths[j])
filename_without_ext = os.path.splitext(original_filename)[0]
final_file_name = f"group-{group_index+1}-{j+1}-{filename_without_ext}"
output_path = os.path.join(self.output_folder, f"{final_file_name}.jpg")
log.info(f" Saving processed image to {output_path}")
img.convert('RGB').save(output_path, quality=100)
if self.caption:
self.create_caption_file(source_paths[j], group_index, final_file_name)
def create_caption_file(self, source_path, group_index, caption_filename):
dirpath = os.path.dirname(source_path)
caption = os.path.basename(dirpath).split('_')[-1]
caption_filename = caption_filename + self.caption_ext
caption_path = os.path.join(self.output_folder, caption_filename)
with open(caption_path, 'w') as f:
f.write(caption)
def copy_other_files(self, group, group_index):
for j, path in enumerate(group):
dirpath, original_filename = os.path.split(path)
original_basename, original_ext = os.path.splitext(original_filename)
for filename in os.listdir(dirpath):
if filename.endswith('.npz'): # Skip .npz
continue
basename, ext = os.path.splitext(filename)
if basename == original_basename and ext != original_ext:
shutil.copy2(os.path.join(dirpath, filename), os.path.join(self.output_folder, f"group-{group_index+1}-{j+1}-{filename}"))
def process_images(self):
images = self.get_image_paths()
groups = self.group_images(images)
for i, group in enumerate(groups):
log.info(f"Processing group {i+1} with {len(group)} images...")
self.process_group(group, i)
def process_group(self, group, group_index):
if len(group) > 0:
aspect_ratios = self.get_aspect_ratios(group)
avg_aspect_ratio = np.mean(aspect_ratios)
if self.pad:
padded_images = self.pad_images(group, avg_aspect_ratio)
self.resize_and_save_images(padded_images, group_index, group)
else:
cropped_images = self.crop_images(group, avg_aspect_ratio)
self.resize_and_save_images(cropped_images, group_index, group)
if not self.do_not_copy_other_files:
self.copy_other_files(group, group_index)
def pad_images(self, group, avg_aspect_ratio):
padded_images = []
for j, path in enumerate(group):
with Image.open(path) as img:
log.info(f" Processing image {j+1}: {path}")
img = self.pad_image(img, avg_aspect_ratio)
padded_images.append(img)
return padded_images
def pad_image(self, img, avg_aspect_ratio):
img_aspect_ratio = img.width / img.height
if img_aspect_ratio < avg_aspect_ratio:
# Too tall, increase width
new_width = avg_aspect_ratio * img.height
pad_width = int((new_width - img.width) / 2)
img = ImageOps.expand(img, border=(pad_width, 0), fill='black')
else:
# Too wide, increase height
new_height = img.width / avg_aspect_ratio
pad_height = int((new_height - img.height) / 2)
img = ImageOps.expand(img, border=(0, pad_height), fill='black')
return img
def main():
parser = argparse.ArgumentParser(description='Process groups of images.')
parser.add_argument('input_folder', type=str, help='Input folder containing images')
parser.add_argument('output_folder', type=str, help='Output folder to store processed images')
parser.add_argument('group_size', type=int, help='Number of images in each group')
parser.add_argument('--include_subfolders', action='store_true', help='Include subfolders in search for images')
parser.add_argument('--do_not_copy_other_files', '--no_copy', dest='do_not_copy_other_files', action='store_true', help='Do not copy other files with the same name as images')
parser.add_argument('--pad', action='store_true', help='Pad images instead of cropping them')
parser.add_argument('--caption', action='store_true', help='Create a caption file for each image')
parser.add_argument('--caption_ext', type=str, default='.txt', help='Extension for the caption file')
args = parser.parse_args()
processor = ImageProcessor(args.input_folder, args.output_folder, args.group_size, args.include_subfolders, args.do_not_copy_other_files, args.pad, args.caption, args.caption_ext)
processor.process_images()
if __name__ == "__main__":
main()
|