You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

101 lines
2.5 KiB

import os
from config import Config as Conf
from utils import check_shape
from transform.gan.mask import CorrectToMask, MaskrefToMaskdet, MaskfinToNude
from transform.opencv.correct import DressToCorrect, ColorTransfer
from transform.opencv.mask import MaskToMaskref, MaskdetToMaskfin
from transform.opencv.resize import ImageToResized, ImageToCrop, ImageToOverlay, ImageToResizedCrop, ImageToRescale
def shift_step(args, shift_start_add=0, shift_end_add=0):
if not args['steps']:
args['steps'] = (0, 5)
args['steps'] = (
args['steps'][0] + shift_start_add,
args['steps'][1] + shift_end_add
def shift_starting(args):
if args.get('steps') and args['steps'][0] != 0:
shift_step(args, shift_start_add=1)
def shift_ending(args, p):
if args.get('steps') and args['steps'][1] == len(p) - 1:
shift_step(args, shift_end_add=1)
def add_tail(args, p, add):
p = [add] + p
shift_ending(args, p)
return p
def add_head(args, p, add):
p = p + [add]
shift_ending(args, p)
return p
def overlay(args, p):
p = add_tail(args, p, ImageToResized)
p = add_tail(args, p, ImageToCrop)
p = add_head(args, p, ImageToOverlay)
return p
def auto_resize(args, p):
return add_tail(args, p, ImageToResized)
def auto_resize_crop(args, p):
return add_tail(args, p, ImageToResizedCrop)
def auto_rescale(args, p):
return add_tail(args, p, ImageToRescale)
def is_file(args, path):
if not os.path.isfile(path):
return False
for mod in (overlay, auto_resize, auto_resize_crop, auto_rescale):
if args.get(mod.__name__):
return True
if not args['ignore_size']:
Conf.log.warn('Image Size Requirements Unchecked.')
return True
def scale_mod(args, p):
for mod in (overlay, auto_resize, auto_resize_crop, auto_rescale):
if args.get(mod.__name__):
return mod(args, p)
if os.path.isfile(Conf.args["input"]):
is_file(args, Conf.args["input"])
return p
def select_phases(args):
Select the transformation phases to use following args parameters.
:return: <ImageTransform[]> list of image transformation
phases = [DressToCorrect, CorrectToMask, MaskToMaskref,
MaskrefToMaskdet, MaskdetToMaskfin, MaskfinToNude]
phases = scale_mod(args, phases)
if args['color_transfer']:
phases = add_head(args, phases, ColorTransfer)
return phases