File size: 8,608 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 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 |
# This code sorts a collection of images in a given directory by their aspect ratio, groups
# them into batches of a given size, crops each image in a batch to the average aspect ratio
# of that batch, and saves the cropped images in a specified directory. The user provides
# the paths to the input directory and the output directory, as well as the desired batch
# size. The program drops any images that do not fit exactly into the batches.
import os
import cv2
import argparse
import shutil
def aspect_ratio(img_path):
"""Return aspect ratio of an image"""
image = cv2.imread(img_path)
height, width = image.shape[:2]
aspect_ratio = float(width) / float(height)
return aspect_ratio
def sort_images_by_aspect_ratio(path):
"""Sort all images in a folder by aspect ratio"""
images = []
for filename in os.listdir(path):
if filename.endswith(".jpg") or filename.endswith(".jpeg") or filename.endswith(".png") or filename.endswith(".webp"):
print(filename)
img_path = os.path.join(path, filename)
images.append((img_path, aspect_ratio(img_path)))
# sort the list of tuples based on the aspect ratio
sorted_images = sorted(images, key=lambda x: x[1])
return sorted_images
def create_groups(sorted_images, n_groups):
"""Create n groups from sorted list of images"""
n = len(sorted_images)
size = n // n_groups
groups = [sorted_images[i * size : (i + 1) * size] for i in range(n_groups - 1)]
groups.append(sorted_images[(n_groups - 1) * size:])
return groups
def average_aspect_ratio(group):
"""Calculate average aspect ratio for a group"""
aspect_ratios = [aspect_ratio for _, aspect_ratio in group]
avg_aspect_ratio = sum(aspect_ratios) / len(aspect_ratios)
print(f"Average aspect ratio for group: {avg_aspect_ratio}")
return avg_aspect_ratio
def center_crop_image(image, target_aspect_ratio):
"""Crop the input image to the target aspect ratio.
The function calculates the crop region for the input image based on its current aspect ratio and the target aspect ratio.
Args:
image: A numpy array representing the input image.
target_aspect_ratio: A float representing the target aspect ratio.
Returns:
A numpy array representing the cropped image.
"""
height, width = image.shape[:2]
current_aspect_ratio = float(width) / float(height)
if current_aspect_ratio == target_aspect_ratio:
return image
if current_aspect_ratio > target_aspect_ratio:
new_width = int(target_aspect_ratio * height)
x_start = (width - new_width) // 2
cropped_image = image[:, x_start:x_start+new_width]
else:
new_height = int(width / target_aspect_ratio)
y_start = (height - new_height) // 2
cropped_image = image[y_start:y_start+new_height, :]
return cropped_image
def copy_related_files(img_path, save_path):
"""
Copy all files in the same directory as the input image that have the same base name as the input image to the
output directory with the corresponding new filename.
:param img_path: Path to the input image.
:param save_path: Path to the output image.
"""
# Get the base filename and directory
img_dir, img_basename = os.path.split(img_path)
img_base, img_ext = os.path.splitext(img_basename)
save_dir, save_basename = os.path.split(save_path)
save_base, save_ext = os.path.splitext(save_basename)
# Create the output directory if it does not exist
if not os.path.exists(save_dir):
os.makedirs(save_dir)
# Loop over all files in the same directory as the input image
try:
for filename in os.listdir(img_dir):
# Skip files with the same name as the input image
if filename == img_basename:
continue
# Check if the file has the same base name as the input image
file_base, file_ext = os.path.splitext(filename)
if file_base == img_base:
# Build the new filename and copy the file
new_filename = os.path.join(save_dir, f"{save_base}{file_ext}")
shutil.copy2(os.path.join(img_dir, filename), new_filename)
except OSError as e:
print(f"Error: {e}") # Handle errors from os.listdir()
def save_resized_cropped_images(group, folder_name, group_number, avg_aspect_ratio, use_original_name=False):
"""Crop and resize all images in the input group to the smallest resolution, and save them to a folder.
Args:
group: A list of tuples, where each tuple contains the path to an image and its aspect ratio.
folder_name: A string representing the name of the folder to save the images to.
group_number: An integer representing the group number.
avg_aspect_ratio: A float representing the average aspect ratio of the images in the group.
use_original_name: A boolean indicating whether to save the images with their original file names.
"""
if not os.path.exists(folder_name):
os.makedirs(folder_name)
# get the smallest size of the images
smallest_res = float("inf")
for img_path, _ in group:
image = cv2.imread(img_path)
cropped_image = center_crop_image(image, avg_aspect_ratio)
height, width = cropped_image.shape[:2]
image_res = height * width
if image_res < smallest_res:
smallest_res = image_res
small_height, small_width = height, width
# resize all images to the smallest resolution of the images in the group
for i, (img_path, aspect_ratio) in enumerate(group):
image = cv2.imread(img_path)
cropped_image = center_crop_image(image, avg_aspect_ratio)
# resized_image = cv2.resize(cropped_image, (small_width, small_height))
if use_original_name:
save_name = os.path.basename(img_path)
else:
save_name = f"group_{group_number}_{i}.jpg"
save_path = os.path.join(folder_name, save_name)
cv2.imwrite(save_path, cropped_image)
# Copy matching files named the same as img_path to
copy_related_files(img_path, save_path)
print(f"Saved {save_name} to {folder_name}")
def main():
parser = argparse.ArgumentParser(description='Sort images and crop them based on aspect ratio')
parser.add_argument('input_dir', type=str, help='Path to the directory containing images')
parser.add_argument('output_dir', type=str, help='Path to the directory to save the cropped images')
parser.add_argument('batch_size', type=int, help='Size of the batches to create')
parser.add_argument('--use_original_name', action='store_true', help='Whether to use original file names for the saved images')
args = parser.parse_args()
print(f"Sorting images by aspect ratio in {args.input_dir}...")
if not os.path.exists(args.input_dir):
print(f"Error: Input directory does not exist: {args.input_dir}")
return
if not os.path.exists(args.output_dir):
try:
os.makedirs(args.output_dir)
except OSError:
print(f"Error: Failed to create output directory: {args.output_dir}")
return
sorted_images = sort_images_by_aspect_ratio(args.input_dir)
total_images = len(sorted_images)
print(f'Total images: {total_images}')
if args.batch_size <= 0:
print("Error: Batch size must be greater than 0")
return
group_size = total_images // args.batch_size
print(f'Train batch size: {args.batch_size}, image group size: {group_size}')
remainder = total_images % args.batch_size
if remainder != 0:
print(f'Dropping {remainder} images that do not fit in groups...')
sorted_images = sorted_images[:-remainder]
total_images = len(sorted_images)
group_size = total_images // args.batch_size
print('Creating groups...')
groups = create_groups(sorted_images, group_size)
print(f"Created {len(groups)} groups")
print('Saving cropped and resize images...')
for i, group in enumerate(groups):
avg_aspect_ratio = average_aspect_ratio(group)
print(f"Processing group {i+1} with {len(group)} images...")
try:
save_resized_cropped_images(group, args.output_dir, i+1, avg_aspect_ratio, args.use_original_name)
except Exception as e:
print(f"Error: Failed to save images in group {i+1}: {e}")
print('Done')
if __name__ == '__main__':
main() |