dreampower/processing/utils.py

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_starting(args)
    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']:
        check_shape(path)
    else:
        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