Mirrors
/
waifu-diffusion
mirror of https://github.com/harubaru/waifu-diffusion.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add centercropping

This commit is contained in:
harubaru 2022-09-22 19:45:36 -07:00 committed by GitHub
parent 5b169173d5
commit d0ec5be718
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 27 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
danbooru_data/download.py
View File

@ -22,11 +22,12 @@ parser = argparse.ArgumentParser()
parser.add_argument('--file', '-f', type=str, required=False, default='links.json')
parser.add_argument('--out_file', '-o', type=str, required=False, default='dataset-%06d.tar')
parser.add_argument('--max_size', '-m', type=int, required=False, default=4294967296)
parser.add_argument('--threads', '-p', required=False, default=16)
parser.add_argument('--resize', '-r', required=False, default=768)
parser.add_argument('--threads', '-p', required=False, default=16, type=int)
parser.add_argument('--resize', '-r', required=False, default=768, type=int)
args = parser.parse_args()
def resize_image(image: Image, max_size=(768,768)):
def resize_image(image: Image, max_size=(512,512), center_crop=True):
if not center_crop:
image = ImageOps.contain(image, max_size, Image.LANCZOS)
# resize to integer multiple of 64
w, h = image.size
@ -41,6 +42,17 @@ def resize_image(image: Image, max_size=(768,768)):
resized = image.resize((src_w, src_h), resample=Image.LANCZOS)
res = Image.new("RGB", (w, h))
res.paste(resized, box=(w // 2 - src_w // 2, h // 2 - src_h // 2))
else:
if not image.mode == "RGB":
image = image.convert("RGB")
if center_crop:
img = np.array(image).astype(np.uint8)
crop = min(img.shape[0], img.shape[1])
h, w, = img.shape[0], img.shape[1]
img = img[(h - crop) // 2:(h + crop) // 2,
(w - crop) // 2:(w + crop) // 2]
res = Image.fromarray(img)
res = res.resize(max_size, resample=Image.LANCZOS)
return res