Browse Source

Merge pull request #41 from PommeDroid/argv_change

Argv parsing change and refactoring
tags/v1.2.10
deeppppp 2 years ago
parent
commit
b9796521f8
No account linked to committer's email address
13 changed files with 709 additions and 392 deletions
  1. 0
    373
      argv.py
  2. 129
    0
      argv/__init__.py
  3. 79
    0
      argv/checkpoints.py
  4. 28
    0
      argv/common.py
  5. 21
    0
      argv/daemon.py
  6. 30
    0
      argv/gpu_info.py
  7. 345
    0
      argv/run.py
  8. 0
    1
      checkpoints.py
  9. 60
    0
      daemon.py
  10. 6
    8
      gpu_info.py
  11. 5
    5
      main.py
  12. 3
    4
      processing/__init__.py
  13. 3
    1
      requirements.txt

+ 0
- 373
argv.py View File

@@ -1,373 +0,0 @@
import argparse
import copy
import json
import os
import re
import sys
from json import JSONDecodeError

import checkpoints
import gpu_info
from main import main
from config import Config as conf
from utils import cv2_supported_extension, json_to_argv, check_image_file_validity


class ArgvParser:
parser = argparse.ArgumentParser(
description="Dreampower CLI application that allow to transform photos of people for private entertainment",
add_help=False
)

@staticmethod
def config_args(args, json_data=None):
"""
Config, do check for a Namespace and give this dict representation of args
:param json_path: <string> a json config file to use to update args
:param args: <Namespace> args Namespace
:return: <dict> dict representation of args
"""

def config_checkpoints(a):
checkpoints_dir = a.checkpoints
a.checkpoints = {
'correct_to_mask': os.path.join(str(a.checkpoints), "cm.lib"),
'maskref_to_maskdet': os.path.join(str(a.checkpoints), "mm.lib"),
'maskfin_to_nude': os.path.join(str(a.checkpoints), "mn.lib"),
}
for _, v in a.checkpoints.items():
if not os.path.isfile(v):
ArgvParser.parser.error(
"Checkpoints file not found in directory {}. "
"You can download them using : {} checkpoints download".format(checkpoints_dir, sys.argv[0])
)

def config_body_parts_prefs(a):
a.prefs = {
"titsize": a.bsize,
"aursize": a.asize,
"nipsize": a.nsize,
"vagsize": a.vsize,
"hairsize": a.hsize
}

def config_gpu_ids(a):
if a.cpu:
a.gpu_ids = None
elif a.gpu:
a.gpu_ids = a.gpu
else:
a.gpu_ids = None if not gpu_info.get_info()['has_cuda'] else [0]

def config_args_in(a):
if not a.input:
ArgvParser.parser.error("-i, --input INPUT is required.")
elif not os.path.isdir(a.input) and not os.path.isfile(a.input):
ArgvParser.parser.error("Input {} file or directory doesn't exist.".format(a.input))
elif os.path.isfile(a.input) and os.path.splitext(a.input)[1] \
not in cv2_supported_extension() + [".gif"]:
ArgvParser.parser.error("Input {} file not supported format.".format(a.input))
if os.path.isfile(a.input):
check_image_file_validity(a.input)

def config_args_out(a):
if os.path.isfile(a.input) and not a.output:
_, extension = os.path.splitext(a.input)
a.output = "output{}".format(extension)
elif a.output and os.path.isfile(a.input) and os.path.splitext(a.output)[1] \
not in cv2_supported_extension() + [".gif"]:
ArgvParser.parser.error("Output {} file not a supported format.".format(a.output))

def config_args_altered(a):
if a.steps and not a.altered:
ArgvParser.parser.error("--steps requires --altered.")
elif a.steps and a.altered:
if not os.path.isdir(a.altered):
ArgvParser.parser.error("{} directory doesn't exist.".format(a.altered))

def config_all(a):
config_checkpoints(a)
config_gpu_ids(a)
config_body_parts_prefs(a)
config_args_in(a)
config_args_out(a)
config_args_altered(a)
return a

def merge_args_json_in_dict(a, json_data=None):
def filter_conflict_args(l1, l2):
# l2 args got priority on l1
l1 = copy.copy(l1)
l2 = copy.copy(l2)
# Handle special cases for ignoring arguments in json file if provided in command line
if "--cpu" in l2 or "--gpu" in l2:
l1 = list(filter(lambda x: x not in ("--cpu", "--gpu"), l1))

if "--auto-resize" in l2 or "--auto-resize-crop" in l2 \
or "--auto-rescale" in l2 or "--overlay" in l2:
l1 = list(filter(lambda x: x not in ("--auto-resize", "--auto-resize-crop", "--auto-rescale"), l1))
if "--overlay" in l1:
del l1[l1.index("--overlay"):l1.index("--overlay") + 1]

return l1 + l2

cmdline_args = []
if not json_data and not a.json_args:
return vars(a)
elif json_data and a.json_args:
cmdline_args = filter_conflict_args(json_to_argv(json_data), json_to_argv(a.json_args))
elif json_data and not a.json_args:
cmdline_args = json_to_argv(json_data)
elif not json_data and a.json_args:
cmdline_args = json_to_argv(a.json_args)

cmdline_args = filter_conflict_args(cmdline_args, sys.argv[1:])

i = 0
while i < len(cmdline_args):
if "--json-args" == cmdline_args[i]:
del cmdline_args[i:i + 1]
i += 1

return vars(config_all(ArgvParser.parser.parse_args(cmdline_args)))

args = copy.deepcopy(args)
if args.func == main:
config_all(args)
return merge_args_json_in_dict(args, json_data)
return vars(args)

@staticmethod
def run():
"""
Run argparse for Dreampower
:return: None
"""
ArgvParser.parser.add_argument('-h', '--help', action='help', default=argparse.SUPPRESS,
help='Show this help message and exit.')
subparsers = ArgvParser.parser.add_subparsers()
ArgvParser.parser.add_argument(
"-d", "--debug", action="store_true", help="Enable log debug mod."
)
ArgvParser.parser.add_argument(
"-i", "--input", help="Path of the photo or directory to transform ."
)
ArgvParser.parser.add_argument(
"-o", "--output",
help="Path of the file or the directory where the transformed photo(s)"
"will be saved. Default : output<input extension>",
)
processing_mod = ArgvParser.parser.add_mutually_exclusive_group()
processing_mod.add_argument(
"--cpu",
default=False,
action="store_true",
help="Force photo processing with CPU (slower)",
)
processing_mod.add_argument(
"--gpu",
action="append",
type=int,
help="ID of the GPU to use for processing. "
"It can be used multiple times to specify multiple GPUs "
"(Example: --gpu 0 --gpu 1 --gpu 2). Default : 0"
)
ArgvParser.parser.add_argument(
"--bsize",
type=float,
default=1,
help="Boob size scalar best results 0.3 - 2.0",
)
ArgvParser.parser.add_argument(
"--asize",
type=float,
default=1,
help="Areola size scalar best results 0.3 - 2.0",
)
ArgvParser.parser.add_argument(
"--nsize",
type=float,
default=1,
help="Nipple size scalar best results 0.3 - 2.0",
)
ArgvParser.parser.add_argument(
"--vsize",
type=float,
default=1,
help="Vagina size scalar best results 0.3 - 1.5",
)
ArgvParser.parser.add_argument(
"--hsize",
type=float,
default=0,
help="Pubic hair size scalar best results set to 0 to disable",
)
ArgvParser.parser.add_argument(
"-n", "--n-runs", type=int, default=1, help="Number of times to process input. Default : 1",
)
ArgvParser.parser.add_argument(
"--n-cores", type=int, default=1, help="Number of cpu cores to use. Default : 1",
)

scale_mod = ArgvParser.parser.add_mutually_exclusive_group()
scale_mod.add_argument(
"--auto-resize",
action="store_true",
default=False,
help="Scale and pad image to 512x512 (maintains aspect ratio).",
)
scale_mod.add_argument(
"--auto-resize-crop",
action="store_true",
default=False,
help="Scale and crop image to 512x512 (maintains aspect ratio).",
)
scale_mod.add_argument(
"--auto-rescale",
action="store_true",
default=False,
help="Scale image to 512x512.",
)

def check_crops_coord():
def type_func(a):
if not re.match(r"^\d+,\d+:\d+,\d+$", a):
raise argparse.ArgumentTypeError("Incorrect coordinates format. "
"Valid format is <x_top_left>,"
"<y_top_left>:<x_bot_right>,<x_bot_right>.")
return tuple(int(x) for x in re.findall(r'\d+', a))

return type_func

scale_mod.add_argument(
"--overlay",
type=check_crops_coord(),
help="Processing the part of the image given by the coordinates "
"(<x_top_left>,<y_top_left>:<x_bot_right>,<x_bot_right>) and overlay the result on the original image."
)

scale_mod.add_argument(
"--ignore-size",
action="store_true",
help="Ignore image size checks."
)

def check_steps_args():
def type_func(a):
if not re.match(r"^[0-5]:[0-5]$", a):
raise argparse.ArgumentTypeError("Incorrect skip format. "
"Valid format is <starting step>:<ending step>.\n"
"Steps are : \n"
"0 : dress -> correct [OPENCV]\n"
"1 : correct -> mask [GAN]\n"
"2 : mask -> maskref [OPENCV]\n"
"3 : maskref -> maskdet [GAN]\n"
"4 : maskdet -> maskfin [OPENCV]\n"
"5 : maskfin -> nude [GAN]"
)

steps = tuple(int(x) for x in re.findall(r'\d+', a))

if steps[0] > steps[1]:
raise argparse.ArgumentTypeError("The ending step should be greater than starting the step.")

return steps[0], steps[1] + 1

return type_func

ArgvParser.parser.add_argument(
"--color-transfer",
action="store_true",
help="Transfers the color distribution from the input image to the output image."
)

ArgvParser.parser.add_argument(
"-s",
"--steps",
type=check_steps_args(),
help="Select a range of steps to execute <starting step>:<ending step>."
"Steps are : \n"
"0 : dress -> correct [OPENCV]\n"
"1 : correct -> mask [GAN]\n"
"2 : mask -> maskref [OPENCV]\n"
"3 : maskref -> maskdet [GAN]\n"
"4 : maskdet -> maskfin [OPENCV]\n"
"5 : maskfin -> nude [GAN]"
)

ArgvParser.parser.add_argument(
"-a",
"--altered",
help="Path of the directory where steps images transformation are write."
)

ArgvParser.parser.add_argument(
"-c",
"--checkpoints",
default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "checkpoints"),
help="Path of the directory containing the checkpoints. Default : ./checkpoints"
)

def check_json_args_file():
def type_func(a):
try:
if os.path.isfile(a):
with open(a, 'r') as f:
j = json.load(f)
else:
j = json.loads(str(a))
except JSONDecodeError:
raise argparse.ArgumentTypeError(
"Arguments json {} is not in valid JSON format.".format(a))
return j

return type_func

ArgvParser.parser.add_argument(
"-j",
"--json-args",
type=check_json_args_file(),
help="Load arguments from json files or json string. "
"If a command line argument is also provide the json value will be ignore for this argument.",
)
ArgvParser.parser.add_argument(
"--json-folder-name",
default="settings.json",
help="Path to the json per folder configuration to looks for when processing folder. Default: settings.json"
)

ArgvParser.parser.add_argument(
"-v",
"--version",
action='version', version='%(prog)s {}'.format(conf.version)
)

gpu_info_parser = subparsers.add_parser('gpu-info')
gpu_info_subparser = gpu_info_parser.add_subparsers()
gpu_info_json_parser = gpu_info_subparser.add_parser('json')

checkpoints_parser = subparsers.add_parser('checkpoints')
checkpoints_parser_subparser = checkpoints_parser.add_subparsers()
checkpoints_parser_info_parser = checkpoints_parser_subparser.add_parser('download')

checkpoints_parser.add_argument(
"-v",
"--version",
action='version', version='checkpoints {}'.format(conf.checkpoints_version)
)

# Register Command Handlers
ArgvParser.parser.set_defaults(func=main)
gpu_info_parser.set_defaults(func=gpu_info.main)
gpu_info_json_parser.set_defaults(func=gpu_info.json)
checkpoints_parser.set_defaults(func=checkpoints.main)
checkpoints_parser_info_parser.set_defaults(func=checkpoints.download)

# Show usage is no args is provided
if len(sys.argv) == 1:
ArgvParser.parser.print_usage()
ArgvParser.parser.exit()

args = ArgvParser.parser.parse_args()
conf.args = ArgvParser.config_args(args)
args.func(args)

+ 129
- 0
argv/__init__.py View File

@@ -0,0 +1,129 @@
import argparse
import copy
import logging
import sys
from config import Config as conf
from argv.checkpoints import init_checkpoints_sub_parser, check_args_checkpoints_parser, set_args_checkpoints_parser
from argv.common import arg_help, arg_debug, arg_version
from argv.daemon import init_daemon_sub_parser
from argv.gpu_info import init_gpu_info_sub_parser
from argv.run import init_run_parser, check_args_run_parser, set_args_run_parser
from utils import setup_log, json_to_argv


class Parser:
parser = None


def run():
Parser.parser = init_parser()

if len(sys.argv) == 1:
Parser.parser.print_usage()
Parser.parser.exit()

args = Parser.parser.parse_args()

conf.log = setup_log(logging.DEBUG) if args.debug else setup_log()
args = config_args(Parser.parser, args)

conf.log.debug(args)

conf.args = vars(args)
args.func(args)


def init_parser():
parser = argparse.ArgumentParser(
description="Dreampower CLI application that allow to transform photos of people for private entertainment",
add_help=False
)

# add args main
arg_help(parser)
arg_debug(parser)
arg_version(parser)

# add subparsers
subparsers = parser.add_subparsers(dest="mode")
init_run_parser(subparsers)
init_checkpoints_sub_parser(subparsers)
init_gpu_info_sub_parser(subparsers)
init_daemon_sub_parser(subparsers)

return parser


def config_args(parser, args, json_data=None):
args = copy.deepcopy(args)
set_args_parser(parser, args)
check_args_parser(parser, args)
if args.mode == "run":
return merge_json_args(parser, args, json_data)
else:
return args


def check_args_parser(parser, args):
if args.mode == "run":
check_args_run_parser(parser, args)

if args.mode == "checkpoints":
check_args_checkpoints_parser(parser, args)

return parser


def set_args_parser(parser, args):
if args.mode == "run":
set_args_run_parser(args)

if args.mode == "checkpoints":
set_args_checkpoints_parser(args)

return parser


def merge_json_args(parser, args, json_data=None):
def filter_conflict_args(l1, l2):
# l2 args got priority on l1
l1 = copy.copy(l1)
l2 = copy.copy(l2)
# Handle special cases for ignoring arguments in json file if provided in command line
if "--cpu" in l2 or "--gpu" in l2:
l1 = list(filter(lambda x: x not in ("--cpu", "--gpu"), l1))

if "--auto-resize" in l2 or "--auto-resize-crop" in l2 \
or "--auto-rescale" in l2 or "--overlay" in l2:
l1 = list(filter(lambda x: x not in ("--auto-resize", "--auto-resize-crop", "--auto-rescale"), l1))
if "--overlay" in l1:
del l1[l1.index("--overlay"):l1.index("--overlay") + 1]

return l1 + l2

# merge args
cmdline_args = []
if not json_data and not args.json_args:
return args
elif json_data and args.json_args:
cmdline_args = filter_conflict_args(json_to_argv(json_data), json_to_argv(args.json_args))
elif json_data and not args.json_args:
cmdline_args = json_to_argv(json_data)
elif not json_data and args.json_args:
cmdline_args = json_to_argv(args.json_args)

cmdline_args = filter_conflict_args(cmdline_args, sys.argv[1:])

# remove json-args arg
i = 0
while i < len(cmdline_args):
if "--json-args" == cmdline_args[i]:
del cmdline_args[i:i + 2]
i += 1

# replace run arg
cmdline_args.index('run')
cmdline_args.remove('run')
cmdline_args.insert(0, 'run')

return config_args(parser, parser.parse_args(cmdline_args))

+ 79
- 0
argv/checkpoints.py View File

@@ -0,0 +1,79 @@
import os
import sys

import checkpoints
from config import Config as conf
from argv.common import arg_help, arg_debug


def init_checkpoints_sub_parser(subparsers):
checkpoints_parser = subparsers.add_parser(
'checkpoints',
description="Handle checkpoints for dreampower.",
help="Handle checkpoints for dreampower.",
add_help=False
)

# add checkpoints arguments
arg_checkpoints(checkpoints_parser)

arg_help(checkpoints_parser)
arg_debug(checkpoints_parser)
arg_version(checkpoints_parser)

# add download subparser
checkpoints_parser_subparser = checkpoints_parser.add_subparsers()
checkpoints_parser_info_parser = checkpoints_parser_subparser.add_parser(
'download',
description="Download checkpoints for dreampower.",
help="Download checkpoints for dreampower."
)

checkpoints_parser.set_defaults(func=checkpoints.main)
checkpoints_parser_info_parser.set_defaults(func=checkpoints.download)

return checkpoints_parser


def set_args_checkpoints_parser(args):
set_arg_checkpoints(args)


def check_args_checkpoints_parser(parser, args):
check_arg_checkpoints(parser, args)


def check_arg_checkpoints(parser, args):
conf.log.debug(args.checkpoints)
for _, v in args.checkpoints.items():
if not os.path.isfile(v):
parser.error(
"Checkpoints file not found. "
"You can download them using : {} checkpoints download".format(sys.argv[0])
)


def set_arg_checkpoints(args):
conf.log.debug(args.checkpoints)
args.checkpoints = {
'correct_to_mask': os.path.join(str(args.checkpoints), "cm.lib"),
'maskref_to_maskdet': os.path.join(str(args.checkpoints), "mm.lib"),
'maskfin_to_nude': os.path.join(str(args.checkpoints), "mn.lib"),
}


def arg_checkpoints(parser):
parser.add_argument(
"-c",
"--checkpoints",
default=os.path.join(os.getcwd(), "checkpoints"),
help="Path of the directory containing the checkpoints. Default : ./checkpoints"
)


def arg_version(parser):
parser.add_argument(
"-v",
"--version",
action='version', version='checkpoints {}'.format(conf.checkpoints_version)
)

+ 28
- 0
argv/common.py View File

@@ -0,0 +1,28 @@
from config import Config as conf


def arg_debug(parser):
parser.add_argument(
"-d",
"--debug",
action="store_true",
help="Enable log debug mod."
)


def arg_help(parser):
parser.add_argument(
'-h',
'--help',
action='help',
help='Show this help message and exit.'
)


def arg_version(parser):
parser.add_argument(
"-v",
"--version",
action='version', version='%(prog)s {}'.format(conf.version)
)


+ 21
- 0
argv/daemon.py View File

@@ -0,0 +1,21 @@
import daemon
from argv.common import arg_debug, arg_help, arg_version


def init_daemon_sub_parser(subparsers):
daemon_parser = subparsers.add_parser(
'daemon',
description="Running dreampower on daemon mode.",
help="Running dreampower on daemon mode.",
add_help=False
)
daemon_parser.set_defaults(func=daemon.main)



# add daemon arguments
arg_help(daemon_parser)
arg_debug(daemon_parser)
arg_version(daemon_parser)

return daemon_parser

+ 30
- 0
argv/gpu_info.py View File

@@ -0,0 +1,30 @@
import gpu_info
from argv.common import arg_debug, arg_help, arg_version


def init_gpu_info_sub_parser(subparsers):
gpu_info_parser = subparsers.add_parser(
'gpu-info',
description="Getting GPU capabilities information for processing with dreampower",
help="Getting GPU capabilities information for processing with dreampower",
add_help=False
)
gpu_info_parser.set_defaults(func=gpu_info.main)

# add gpu-info arguments
arg_json(gpu_info_parser)

arg_help(gpu_info_parser)
arg_debug(gpu_info_parser)
arg_version(gpu_info_parser)

return gpu_info_parser


def arg_json(parser):
parser.add_argument(
"-j",
"--json",
action='store_true',
help=""
)

+ 345
- 0
argv/run.py View File

@@ -0,0 +1,345 @@
import json
import os
import re
from json import JSONDecodeError

import gpu_info
import main
from argv.checkpoints import check_arg_checkpoints, set_arg_checkpoints, arg_checkpoints
from argv.common import arg_debug, arg_help, arg_version
from utils import check_image_file_validity, cv2_supported_extension


def init_run_parser(subparsers):
run_parser = subparsers.add_parser(
'run',
description="Process image(s) with dreampower.",
help="Process image(s) with dreampower.",
add_help=False
)
run_parser.set_defaults(func=main.main)

# conflicts handler
processing_mod = run_parser.add_mutually_exclusive_group()
scale_mod = run_parser.add_mutually_exclusive_group()

# add run arguments
arg_input(run_parser)
arg_output(run_parser)

arg_auto_rescale(scale_mod)
arg_auto_resize(scale_mod)
arg_auto_resize_crop(scale_mod)
arg_overlay(scale_mod)
arg_ignore_size(scale_mod)

arg_color_transfer(run_parser)

arg_preferences(run_parser)
arg_n_run(run_parser)
arg_step(run_parser)
arg_altered(run_parser)

arg_cpu(processing_mod)
arg_gpu(processing_mod)
arg_checkpoints(run_parser)
arg_n_core(run_parser)

arg_json_args(run_parser)
arg_json_folder_name(run_parser)

arg_help(run_parser)
arg_debug(run_parser)
arg_version(run_parser)


def set_args_run_parser(args):
set_arg_checkpoints(args)
set_arg_preference(args)
set_gpu_ids(args)


def check_args_run_parser(parser, args):
check_arg_input(parser, args)
check_arg_output(parser, args)
check_args_altered(parser, args)
check_arg_checkpoints(parser, args)


def check_args_altered(parser, args):
if args.steps and not args.altered:
parser.error("--steps requires --altered.")
elif args.steps and args.altered:
if not os.path.isdir(args.altered):
parser.error("{} directory doesn't exist.".format(args.altered))


def arg_altered(parser):
parser.add_argument(
"-a",
"--altered",
help="Path of the directory where steps images transformation are write."
)


def arg_auto_rescale(parser):
parser.add_argument(
"--auto-rescale",
action="store_true",
help="Scale image to 512x512.",
)


def arg_auto_resize(parser):
parser.add_argument(
"--auto-resize",
action="store_true",
help="Scale and pad image to 512x512 (maintains aspect ratio).",
)


def arg_auto_resize_crop(parser):
parser.add_argument(
"--auto-resize-crop",
action="store_true",
help="Scale and crop image to 512x512 (maintains aspect ratio).",
)


def arg_color_transfer(parser):
parser.add_argument(
"--color-transfer",
action="store_true",
help="Transfers the color distribution from the input image to the output image."
)


def set_gpu_ids(args):
if args.cpu:
args.gpu_ids = None
elif args.gpu:
args.gpu_ids = args.gpu
else:
args.gpu_ids = None if not gpu_info.get_info()['has_cuda'] else [0]


def arg_cpu(parser):
parser.add_argument(
"--cpu",
action="store_true",
help="Force photo processing with CPU (slower)",
)


def arg_gpu(parser):
parser.add_argument(
"--gpu",
action="append",
type=int,
help="ID of the GPU to use for processing. "
"It can be used multiple times to specify multiple GPUs "
"(Example: --gpu 0 --gpu 1 --gpu 2). Default : 0"
)


def arg_ignore_size(parser):
parser.add_argument(
"--ignore-size",
action="store_true",
help="Ignore image size checks."
)


def arg_input(parser):
parser.add_argument(
"-i",
"--input",
help="Path of the photo or directory to transform .",
)


def check_arg_input(parser, args):
if not args.input:
parser.error("-i, --input INPUT is required.")
if not os.path.isdir(args.input) and not os.path.isfile(args.input):
parser.ArgumentTypeError("Input {} file or directory doesn't exist.".format(args.input))
elif os.path.isfile(args.input) and os.path.splitext(args.input)[1] not in cv2_supported_extension() + [".gif"]:
parser.ArgumentTypeError("Input {} file not supported format.".format(args.input))
if os.path.isfile(args.input):
check_image_file_validity(args.input)
return args.input


def arg_json_args(parser):
def check_json_args_file():
def type_func(a):
try:
if os.path.isfile(a):
with open(a, 'r') as f:
j = json.load(f)
else:
j = json.loads(str(a))
except JSONDecodeError:
raise parser.error(
"Arguments json {} is not in valid JSON format.".format(a))
return j

return type_func

parser.add_argument(
"-j",
"--json-args",
type=check_json_args_file(),
help="Load arguments from json files or json string. "
"If a command line argument is also provide the json value will be ignore for this argument.",
)


def arg_json_folder_name(parser):
parser.add_argument(
"--json-folder-name",
default="settings.json",
help="Path to the json per folder configuration to looks for when processing folder. Default: settings.json"
)


def arg_n_core(parser):
parser.add_argument(
"--n-cores",
type=int,
default=1,
help="Number of cpu cores to use. Default : 1",
)


def arg_n_run(parser):
parser.add_argument(
"-n",
"--n-runs",
type=int,
default=1,
help="Number of times to process input. Default : 1"
)


def arg_output(parser):
parser.add_argument(
"-o",
"--output",
help="Path of the file or the directory where the transformed photo(s) "
"will be saved. Default : output<input extension>"
)


def check_arg_output(parser, args):
if os.path.isfile(args.input) and not args.output:
_, extension = os.path.splitext(args.input)
args.output = "output{}".format(extension)
elif args.output and os.path.isfile(args.input) and os.path.splitext(args.output)[1] \
not in cv2_supported_extension() + [".gif"]:
parser.error("Output {} file not a supported format.".format(args.output))


def arg_overlay(parser):
def check_crops_coord():
def type_func(a):
if not re.match(r"^\d+,\d+:\d+,\d+$", a):
raise parser.error("Incorrect coordinates format. "
"Valid format is <x_top_left>,"
"<y_top_left>:<x_bot_right>,<x_bot_right>.")
return tuple(int(x) for x in re.findall(r'\d+', a))

return type_func

parser.add_argument(
"--overlay",
type=check_crops_coord(),
help="Processing the part of the image given by the coordinates "
"(<x_top_left>,<y_top_left>:<x_bot_right>,<x_bot_right>) and overlay the result on the original image."
)


def set_arg_preference(args):
args.prefs = {
"titsize": args.bsize,
"aursize": args.asize,
"nipsize": args.nsize,
"vagsize": args.vsize,
"hairsize": args.hsize
}


def arg_preferences(parser):
parser.add_argument(
"--bsize",
type=float,
default=1,
help="Boob size scalar best results 0.3 - 2.0",
)

parser.add_argument(
"--asize",
type=float,
default=1,
help="Areola size scalar best results 0.3 - 2.0",
)

parser.add_argument(
"--nsize",
type=float,
default=1,
help="Nipple size scalar best results 0.3 - 2.0",
)

parser.add_argument(
"--vsize",
type=float,
default=1,
help="Vagina size scalar best results 0.3 - 1.5",
)

parser.add_argument(
"--hsize",
type=float,
default=0,
help="Pubic hair size scalar best results set to 0 to disable",
)


def arg_step(parser):
def check_steps_args():
def type_func(a):
if not re.match(r"^[0-5]:[0-5]$", a):
raise parser.error("Incorrect skip format. "
"Valid format is <starting step>:<ending step>.\n"
"Steps are : \n"
"0 : dress -> correct [OPENCV]\n"
"1 : correct -> mask [GAN]\n"
"2 : mask -> maskref [OPENCV]\n"
"3 : maskref -> maskdet [GAN]\n"
"4 : maskdet -> maskfin [OPENCV]\n"
"5 : maskfin -> nude [GAN]"
)

steps = tuple(int(x) for x in re.findall(r'\d+', a))

if steps[0] > steps[1]:
raise parser.error("The ending step should be greater than starting the step.")

return steps[0], steps[1] + 1

return type_func

parser.add_argument(
"-s",
"--steps",
type=check_steps_args(),
help="Select a range of steps to execute <starting step>:<ending step>."
"Steps are : \n"
"0 : dress -> correct [OPENCV]\n"
"1 : correct -> mask [GAN]\n"
"2 : mask -> maskref [OPENCV]\n"
"3 : maskref -> maskdet [GAN]\n"
"4 : maskdet -> maskfin [OPENCV]\n"
"5 : maskfin -> nude [GAN]"
)

+ 0
- 1
checkpoints.py View File

@@ -9,7 +9,6 @@ from utils import setup_log, dl_file, unzip


def main(_):
conf.log = setup_log(logging.DEBUG) if conf.args['debug'] else setup_log()
if sum([1 for x in ["cm.lib", "mm.lib", "mn.lib"] if os.path.isfile(os.path.join(conf.args['checkpoints'], x))]):
conf.log.info("Checkpoints Found In {}".format(conf.args['checkpoints']))
else:

+ 60
- 0
daemon.py View File

@@ -0,0 +1,60 @@
import os
import sys
import time

from watchdog.events import FileSystemEventHandler
from watchdog.observers import Observer

from config import Config as conf
from transform.gan.mask import CorrectToMask, MaskrefToMaskdet, MaskfinToNude
from transform.opencv.correct import DressToCorrect
from transform.opencv.mask import MaskToMaskref, MaskdetToMaskfin


class Watcher:
def __init__(self, watching_dir, out_dir):
self.__observer = Observer()
self.__watching_dir = watching_dir
self.__out_dir = out_dir

if not os.path.isdir(self.__watching_dir):
conf.log.error("{} Watching Dir Doesn't Exit.".format(self.__watching_dir))
sys.exit(0)

if not os.path.isdir(self.__out_dir):
conf.log.error("{} Output Dir Doesn't Exit.".format(self.__watching_dir))
sys.exit(0)

def run(self):
event_handler = Handler(self.__out_dir)
self.__observer.schedule(event_handler, self.__watching_dir, recursive=True)
self.__observer.start()
try:
while True:
time.sleep(5)
except KeyboardInterrupt:
self.__observer.stop()
except Exception as e:
self.__observer.stop()
conf.log.error(e)
conf.log.error("An Unhandled Error Occurred.")
sys.exit(1)
self.__observer.join()


class Handler(FileSystemEventHandler):
def __init__(self, out_dir):
self.__out_dir = out_dir
self.__phases = [
DressToCorrect, CorrectToMask, MaskToMaskref, MaskrefToMaskdet, MaskdetToMaskfin, MaskfinToNude
]

def on_created(self, event):
if event.is_directory:
conf.log.debug("Received directory created event {}.".format(event.src_path))
# TODO Implements this


def main():
Watcher("test_dir", "out_dir").run()


+ 6
- 8
gpu_info.py View File

@@ -14,12 +14,10 @@ def get_info():


def main(_):
conf.log = setup_log(logging.DEBUG) if conf.args['debug'] else setup_log()
info = get_info()
conf.log.info("Has Cuda: {}".format(info["has_cuda"]))
for (i, device) in enumerate(info["devices"]):
conf.log.info("GPU {}: {}".format(i, device))


def json(_):
print(j.dumps(get_info()))
if not conf.args['json']:
conf.log.info("Has Cuda: {}".format(info["has_cuda"]))
for (i, device) in enumerate(info["devices"]):
conf.log.info("GPU {}: {}".format(i, device))
else:
print(j.dumps(info))

+ 5
- 5
main.py View File

@@ -1,12 +1,13 @@
import logging
import os
import sys
import time
from multiprocessing import freeze_support

import colorama

import argv
from config import Config as conf
from utils import setup_log, check_shape
from utils import check_shape

from processing import SimpleTransform, FolderImageTransform, MultipleImageTransform
from transform.gan.mask import CorrectToMask, MaskrefToMaskdet, MaskfinToNude
@@ -19,8 +20,6 @@ def main(_):
"""
Main logic entry point
"""
conf.log = setup_log(logging.DEBUG) if conf.args['debug'] else setup_log()
conf.log.debug("Args : {}".format(conf.args))
conf.log.info("Welcome to DreamPower")

if conf.args['gpu_ids']:
@@ -140,5 +139,6 @@ def processing_image_folder(phases):


if __name__ == "__main__":
colorama.init()
freeze_support()
argv.ArgvParser.run()
argv.run()

+ 3
- 4
processing/__init__.py View File

@@ -10,8 +10,6 @@ from multiprocessing.pool import ThreadPool

import cv2
import imageio

import argv
from config import Config as conf
from utils import camel_case_to_str, cv2_supported_extension, read_image, write_image, json_to_argv, check_shape

@@ -130,7 +128,7 @@ class ImageTransform(Process):
conf.log.debug("All Phases : {}".format(self.__phases))
conf.log.debug("To Be Executed Phases : {}".format(self.__phases[self.__starting_step:self.__ending_step]))

path = self.__altered_path if os.path.isfile(self._args['input']) or not self._args.get('folder_altered') \
path = self.__altered_path if os.path.isfile(input_path) or not self._args.get('folder_altered') \
else os.path.join(self._args['folder_altered'], os.path.basename(self.__output_path))

self.__image_steps = [input_path] + [
@@ -292,7 +290,8 @@ class FolderImageTransform(MultipleImageTransform):
.format(folder_path))
return add_folder_altered(self._args)
try:
a = argv.ArgvParser.config_args(argv.ArgvParser.parser.parse_args(sys.argv[1:]), json_data=json_data)
from argv import Parser, config_args
a = config_args(Parser.parser, Parser.parser.parse_args(sys.argv[1:]), json_data=json_data)
conf.log.info("Using {} Configuration for processing {} folder. "
.format(json_path, folder_path))
return add_folder_altered(a)

+ 3
- 1
requirements.txt View File

@@ -6,4 +6,6 @@ torchvision==0.2.2.post3
torch==1.1.0
imageio==2.5.0
coloredlogs==10.0
requests==2.22.0
colorama==0.4.1
requests==2.22.0
watchdog==0.9.0

Loading…
Cancel
Save