Commit 9681419e authored by Milly's avatar Milly

train: fixed preprocess image ratio

parent 7f8ab1ee
import os import os
from PIL import Image, ImageOps from PIL import Image, ImageOps
import math
import platform import platform
import sys import sys
import tqdm import tqdm
...@@ -38,6 +39,8 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro ...@@ -38,6 +39,8 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro
height = process_height height = process_height
src = os.path.abspath(process_src) src = os.path.abspath(process_src)
dst = os.path.abspath(process_dst) dst = os.path.abspath(process_dst)
split_threshold = 0.5
overlap_ratio = 0.2
assert src != dst, 'same directory specified as source and destination' assert src != dst, 'same directory specified as source and destination'
...@@ -78,6 +81,29 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro ...@@ -78,6 +81,29 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro
if process_flip: if process_flip:
save_pic_with_caption(ImageOps.mirror(image), index) save_pic_with_caption(ImageOps.mirror(image), index)
def split_pic(image, inverse_xy):
if inverse_xy:
from_w, from_h = image.height, image.width
to_w, to_h = height, width
else:
from_w, from_h = image.width, image.height
to_w, to_h = width, height
h = from_h * to_w // from_w
if inverse_xy:
image = image.resize((h, to_w))
else:
image = image.resize((to_w, h))
split_count = math.ceil((h - to_h * overlap_ratio) / (to_h * (1.0 - overlap_ratio)))
y_step = (h - to_h) / (split_count - 1)
for i in range(split_count):
y = int(y_step * i)
if inverse_xy:
splitted = image.crop((y, 0, y + to_h, to_w))
else:
splitted = image.crop((0, y, to_w, y + to_h))
yield splitted
for index, imagefile in enumerate(tqdm.tqdm(files)): for index, imagefile in enumerate(tqdm.tqdm(files)):
subindex = [0] subindex = [0]
filename = os.path.join(src, imagefile) filename = os.path.join(src, imagefile)
...@@ -89,26 +115,16 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro ...@@ -89,26 +115,16 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro
if shared.state.interrupted: if shared.state.interrupted:
break break
ratio = img.height / img.width if img.height > img.width:
is_tall = ratio > 1.35 ratio = (img.width * height) / (img.height * width)
is_wide = ratio < 1 / 1.35 inverse_xy = False
else:
if process_split and is_tall: ratio = (img.height * width) / (img.width * height)
img = img.resize((width, height * img.height // img.width)) inverse_xy = True
top = img.crop((0, 0, width, height))
save_pic(top, index)
bot = img.crop((0, img.height - height, width, img.height))
save_pic(bot, index)
elif process_split and is_wide:
img = img.resize((width * img.width // img.height, height))
left = img.crop((0, 0, width, height))
save_pic(left, index)
right = img.crop((img.width - width, 0, img.width, height)) if process_split and ratio < 1.0 and ratio <= split_threshold:
save_pic(right, index) for splitted in split_pic(img, inverse_xy):
save_pic(splitted, index)
else: else:
img = images.resize_image(1, img, width, height) img = images.resize_image(1, img, width, height)
save_pic(img, index) save_pic(img, index)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment