├── LICENSE ├── README.md ├── .gitignore └── resize_and_crop.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Mario Namtao Shianti Larcher 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Resize and Crop Images for Bucketing 2 | 3 | This Python script is designed to resize and crop images for aspect ratio bucketing. It can be particularly useful when working with [Kohya's Stable Diffusion scripts](https://github.com/kohya-ss/sd-scripts) or similar scripts, helping to prevent unexpected assignment to buckets (see [Issue #731](https://github.com/kohya-ss/sd-scripts/issues/731)). 4 | 5 | Given an input image repository, this script resizes and crops images to match the dimensions of the bucket with the closest aspect ratio. Please note the following: 6 | 7 | - Buckets with a width or height greater than that of the image are skipped. 8 | - If no bucket satisfies the constraints, the image is discarded and a warning is raised. 9 | 10 | ## Requirements 11 | 12 | - Python 3.x 13 | - Pillow (Python Imaging Library Fork) 14 | 15 | To install Pillow, run: 16 | ```bash 17 | pip install Pillow 18 | ``` 19 | 20 | ## Usage 21 | 22 | First, clone this repository or download the script. Then, run the script using the command-line interface. 23 | 24 | ### Command-Line Arguments: 25 | 26 | - `--max_sqrt_area`: Maximum square root value of each resolution area. This value is squared to determine the maximum area of the buckets. Default is 1024. 27 | - `--min_size`: Minimum size (either width or height) for the bucket resolutions. Default is 512. 28 | - `--max_size`: Maximum size (either width or height) for the bucket resolutions. Default is 2048. 29 | - `--divisible_by`: Factor by which the bucket widths and heights should be divisible. Default is 64. 30 | - `--input_dir`: Directory containing the input images. 31 | - `--output_dir`: Directory where the resized and cropped images will be saved. 32 | 33 | ### Example Command: 34 | 35 | ```bash 36 | python resize_and_crop.py --input_dir=./input_images --output_dir=./output_images 37 | ``` 38 | 39 | This command will process images from `./input_images`, resize and crop them according to the generated bucket resolutions, and save the output to `./output_images`. 40 | 41 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/#use-with-ide 110 | .pdm.toml 111 | 112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 113 | __pypackages__/ 114 | 115 | # Celery stuff 116 | celerybeat-schedule 117 | celerybeat.pid 118 | 119 | # SageMath parsed files 120 | *.sage.py 121 | 122 | # Environments 123 | .env 124 | .venv 125 | env/ 126 | venv/ 127 | ENV/ 128 | env.bak/ 129 | venv.bak/ 130 | 131 | # Spyder project settings 132 | .spyderproject 133 | .spyproject 134 | 135 | # Rope project settings 136 | .ropeproject 137 | 138 | # mkdocs documentation 139 | /site 140 | 141 | # mypy 142 | .mypy_cache/ 143 | .dmypy.json 144 | dmypy.json 145 | 146 | # Pyre type checker 147 | .pyre/ 148 | 149 | # pytype static type analyzer 150 | .pytype/ 151 | 152 | # Cython debug symbols 153 | cython_debug/ 154 | 155 | # PyCharm 156 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 157 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 158 | # and can be added to the global gitignore or merged into this file. For a more nuclear 159 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 160 | #.idea/ 161 | -------------------------------------------------------------------------------- /resize_and_crop.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import math 3 | import warnings 4 | from pathlib import Path 5 | 6 | from PIL import Image 7 | 8 | 9 | def make_bucket_resolutions( 10 | max_sqrt_area=1024, min_size=512, max_size=2048, divisible_by=64 11 | ): 12 | """ 13 | Generate bucket resolutions based on specified constraints. 14 | 15 | Args: 16 | max_sqrt_area (int): The maximum square root value of each resolution area. 17 | This value is squared to determine the maximum area of the buckets. 18 | min_size (int): The minimum size (either width or height) for the bucket resolutions. 19 | max_size (int): The maximum size (either width or height) for the bucket resolutions. 20 | divisible_by (int): The factor by which the bucket widths and heights should be divisible. 21 | 22 | Returns: 23 | list of tuples: A sorted list of tuples, each containing (aspect_ratio, width, height). 24 | """ 25 | resolutions = set() 26 | 27 | max_sqrt_area = max_sqrt_area // divisible_by 28 | max_area = max_sqrt_area**2 29 | size = max_sqrt_area * divisible_by 30 | resolutions.add((1.0, size, size)) 31 | 32 | size = min_size 33 | while size <= max_size: 34 | width = size 35 | height = min(max_size, (max_area // (width // divisible_by)) * divisible_by) 36 | resolutions.add((width / height, width, height)) 37 | resolutions.add((height / width, height, width)) 38 | size += divisible_by 39 | 40 | resolutions = list(resolutions) 41 | resolutions.sort() 42 | return resolutions 43 | 44 | 45 | def resize_and_crop_images(input_dir, output_dir, bucket_resolutions): 46 | """ 47 | Resize and crop images from an input directory and save them to an output directory, 48 | based on provided resolution buckets. 49 | 50 | Args: 51 | input_dir (str): Path to the directory containing input images. 52 | output_dir (str): Path to the directory where resized and cropped images will be saved. 53 | bucket_resolutions (list of tuples): A sorted list of tuples, each containing (aspect_ratio, width, height). 54 | 55 | Returns: 56 | None: The function saves the processed images to the output directory and does not return a value. 57 | """ 58 | input_dir = Path(input_dir) 59 | output_dir = Path(output_dir) 60 | if not output_dir.exists(): 61 | output_dir.mkdir(parents=True, exist_ok=True) 62 | 63 | # Find all image files in the input directory 64 | image_extensions = ["*.png", "*.jpg", "*.jpeg"] 65 | image_paths = [ 66 | image_path for ext in image_extensions for image_path in input_dir.glob(ext) 67 | ] 68 | 69 | for image_path in image_paths: 70 | image = Image.open(image_path) 71 | width, height = image.size 72 | aspect_ratio = width / height 73 | closest_aspect_ratio = None 74 | target_width = None 75 | target_height = None 76 | closest_aspect_ratio_diff = float("inf") 77 | 78 | for bucket_aspect_ratio, bucket_width, bucket_height in bucket_resolutions: 79 | if width >= bucket_width and height >= bucket_height: 80 | aspect_ratio_diff = abs(bucket_aspect_ratio - aspect_ratio) 81 | if aspect_ratio_diff < closest_aspect_ratio_diff: 82 | closest_aspect_ratio_diff = aspect_ratio_diff 83 | closest_aspect_ratio = bucket_aspect_ratio 84 | target_width = bucket_width 85 | target_height = bucket_height 86 | 87 | if closest_aspect_ratio is not None: 88 | if closest_aspect_ratio > aspect_ratio: 89 | # Resize the image so width matches target_width 90 | new_width = target_width 91 | new_height = math.ceil(height * new_width / width) 92 | image = image.resize((new_width, new_height), Image.LANCZOS) 93 | 94 | # Crop the image so that height matches target_height 95 | if new_height > target_height: 96 | top = (new_height - target_height) // 2 97 | bottom = top + target_height 98 | image = image.crop((0, top, new_width, bottom)) 99 | else: 100 | # Resize the image so height matches target_height 101 | new_height = target_height 102 | new_width = math.ceil(width * new_height / height) 103 | image = image.resize((new_width, new_height), Image.LANCZOS) 104 | 105 | # Crop the image so that width matches target_width 106 | if new_width > target_width: 107 | left = (new_width - target_width) // 2 108 | right = left + target_width 109 | image = image.crop((left, 0, right, new_height)) 110 | 111 | output_path = output_dir / image_path.name 112 | image.save(output_path) 113 | else: 114 | warnings.warn( 115 | f"Skipping {image_path.name}: No suitable aspect ratio bucket found." 116 | ) 117 | 118 | 119 | def parse_args(): 120 | parser = argparse.ArgumentParser( 121 | description="Resize and crop images based on generated bucket resolutions." 122 | ) 123 | 124 | # Arguments for make_bucket_resolutions function 125 | parser.add_argument( 126 | "--max_sqrt_area", 127 | type=int, 128 | default=1024, 129 | help="Maximum square root value of each resolution area. This value is squared to determine the maximum area of the buckets.", 130 | ) 131 | parser.add_argument( 132 | "--min_size", 133 | type=int, 134 | default=512, 135 | help="Minimum size (either width or height) for the bucket resolutions.", 136 | ) 137 | parser.add_argument( 138 | "--max_size", 139 | type=int, 140 | default=2048, 141 | help="Maximum size (either width or height) for the bucket resolutions.", 142 | ) 143 | parser.add_argument( 144 | "--divisible_by", 145 | type=int, 146 | default=64, 147 | help="Factor by which the bucket widths and heights should be divisible.", 148 | ) 149 | 150 | # Arguments for resize_and_crop_images function 151 | parser.add_argument( 152 | "--input_dir", type=str, help="Directory containing input images." 153 | ) 154 | parser.add_argument( 155 | "--output_dir", type=str, help="Directory where output images will be saved." 156 | ) 157 | 158 | return parser.parse_args() 159 | 160 | 161 | if __name__ == "__main__": 162 | args = parse_args() 163 | 164 | bucket_resolutions = make_bucket_resolutions( 165 | args.max_sqrt_area, args.min_size, args.max_size, args.divisible_by 166 | ) 167 | 168 | resize_and_crop_images(args.input_dir, args.output_dir, bucket_resolutions) 169 | --------------------------------------------------------------------------------