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

fix downloader

This commit is contained in:
harubaru 2022-09-22 19:49:08 -07:00
parent d0ec5be718
commit de1470f43b
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
danbooru_data/download.py
View File

@ -17,13 +17,14 @@ from PIL import Image, ImageOps
import argparse import argparse
import shutil import shutil
import numpy as np
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--file', '-f', type=str, required=False, default='links.json') 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('--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('--max_size', '-m', type=int, required=False, default=4294967296)
parser.add_argument('--threads', '-p', required=False, default=16, type=int) parser.add_argument('--threads', '-p', required=False, default=16, type=int)
parser.add_argument('--resize', '-r', required=False, default=768, type=int) parser.add_argument('--resize', '-r', required=False, default=512, type=int)
args = parser.parse_args() args = parser.parse_args()
def resize_image(image: Image, max_size=(512,512), center_crop=True): def resize_image(image: Image, max_size=(512,512), center_crop=True):
@ -45,7 +46,6 @@ def resize_image(image: Image, max_size=(512,512), center_crop=True):
else: else:
if not image.mode == "RGB": if not image.mode == "RGB":
image = image.convert("RGB") image = image.convert("RGB")
if center_crop:
img = np.array(image).astype(np.uint8) img = np.array(image).astype(np.uint8)
crop = min(img.shape[0], img.shape[1]) crop = min(img.shape[0], img.shape[1])
h, w, = img.shape[0], img.shape[1] h, w, = img.shape[0], img.shape[1]