dreamnet
/
dreampower
1
0
Fork 0
Code Issues Pull Requests Releases 11 Wiki Activity
Browse Source

Settings Per Folder for Folder processing #30 

* add json settings per folder
* add autodetect folder / image mod
master
PommeDroid 3 years ago
parent
8f7c67eb1e
commit
40ce0054e0
13 changed files with 579 additions and 417 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 614
      argv.py
  2. 5
      config.py
  3. 22
      main.py
  4. 203
      processing/__init__.py
  5. 2
      requirements.txt
  6. 6
      transform/__init__.py
  7. 8
      transform/gan/__init__.py
  8. 33
      transform/gan/mask.py
  9. 9
      transform/opencv/__init__.py
  10. 22
      transform/opencv/mask.py
  11. 28
      transform/opencv/resize.py
  12. 10
      transform/opencv/watermark.py
  13. 34
      utils.py

614
argv.py
Unescape View File

@ -1,4 +1,5 @@ @@ -1,4 +1,5 @@
import argparse
import copy
import json
import os
import re
@ -8,301 +9,332 @@ from json import JSONDecodeError @@ -8,301 +9,332 @@ from json import JSONDecodeError
import gpu_info
from main import main
from config import Config as conf
from gpu_info import get_info
from utils import cv2_supported_extension
def config_args(parser, args):
def config_checkpoints():
checkpoints = {
'correct_to_mask': os.path.join(args.checkpoints, "cm.lib"),
'maskref_to_maskdet': os.path.join(args.checkpoints, "mm.lib"),
'maskfin_to_nude': os.path.join(args.checkpoints, "mn.lib"),
}
for _, v in checkpoints.items():
if not os.path.isfile(v):
parser.error("Checkpoints file not found in directory {}".format(args.checkpoints))
return checkpoints
def config_body_parts_prefs():
prefs = {
"titsize": args.bsize,
"aursize": args.asize,
"nipsize": args.nsize,
"vagsize": args.vsize,
"hairsize": args.hsize
}
return prefs
def config_gpu_ids():
if args.cpu:
gpu_ids = None
elif args.gpu:
gpu_ids = args.gpu
else:
gpu_ids = None if not gpu_info.get_info()['has_cuda'] else [0]
return gpu_ids
def config_args_in():
if not args.input:
parser.error("-i, --input INPUT is required.")
elif not args.folder and not os.path.isfile(args.input):
parser.error("Input {} file doesn't exist.".format(args.input))
elif args.folder and not os.path.isdir(args.input):
parser.error("Input {} directory doesn't exist.".format(args.input))
elif not args.folder and os.path.splitext(args.input)[1] not in cv2_supported_extension() + [".gif"]:
parser.error("Input {} file not supported format.".format(args.input))
def config_args_out():
if not args.folder and not args.output:
_, extension = os.path.splitext(args.input)
args.output = "output{}".format(extension)
elif args.output and not args.folder and os.path.splitext(args.output)[1] not in cv2_supported_extension():
parser.error("Output {} file not a supported format.".format(args.output))
def config_args_altered():
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.input))
if args.func == main:
conf.args = vars(args)
conf.args['checkpoints'] = config_checkpoints()
conf.args['gpu_ids'] = config_gpu_ids()
conf.args['prefs'] = config_body_parts_prefs()
config_args_in()
config_args_out()
config_args_altered()
def run():
"""
Run argparse for Dreampower
:return: None
"""
parser = argparse.ArgumentParser(
description="Dreampower CLI application that allow to transform photos of people for "
"private entertainment")
subparsers = parser.add_subparsers()
parser.add_argument(
"-d", "--debug", action="store_true", help="enble log debug mod"
)
parser.add_argument(
"-i", "--input", help="path of the photo to transform"
)
parser.add_argument(
"-o",
"--output",
help="path where the transformed photo will be saved. (default: output.<input extension>)",
)
parser.add_argument(
"-f",
"--folder",
action="store_true",
help="Folder mode processing. "
"" # TODO Json config by dir
)
processing_mod = 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)"
" This argument will be ignored if --cpu is active. (default: 0)",
)
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",
)
parser.add_argument(
"-n", "--n_runs", type=int, default=1, help="number of times to process input (default: 1)",
)
parser.add_argument(
"--n_cores", type=int, default=1, help="number of cpu cores to use (default: 1)",
)
scale_mod = 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",
)
from utils import cv2_supported_extension, json_to_argv, check_image_file_validity
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.",
)
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
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]"
)
parser.add_argument(
"-a",
"--altered",
help="path of the directory where steps images transformation are write."
)
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."
class ArgvParser:
parser = argparse.ArgumentParser(
description="Dreampower CLI application that allow to transform photos of people for private entertainment",
add_help=False
)
def check_json_args_file():
def type_func(a):
if not os.path.isfile(a):
raise argparse.ArgumentTypeError(
"Arguments json file {} not found.".format(a))
with open(a) as f:
data = {}
@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):
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 {}".format(a.checkpoints))
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():
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.input))
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 to transform."
)
ArgvParser.parser.add_argument(
"-o",
"--output",
help="Path where the transformed photo 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.",
)
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(
"-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."
)
def check_json_args_file():
def type_func(a):
try:
data = json.load(f)
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 file {} is not in valid JSON format.".format(a))
l = []
for k, v in data.items():
if not isinstance(v, bool):
l.extend(["--{}".format(k), str(v)])
elif v:
l.append("--{}".format(k))
return l
return type_func
parser.add_argument(
"-j",
"--json_args",
type=check_json_args_file(),
help="Load arguments from json files. "
"If a command line argument is also provide the json value will be ignore for this argument.",
)
gpu_info_parser = subparsers.add_parser('gpu-info')
gpu_info_parser.add_argument(
"-j",
"--json",
default=False,
action="store_true",
help="Print GPU info as JSON"
)
# Register Command Handlers
parser.set_defaults(func=main)
gpu_info_parser.set_defaults(func=gpu_info.main)
# Show usage is no args is provided
if len(sys.argv) == 1:
parser.print_usage()
parser.exit()
args = parser.parse_args()
# Handle special cases for ignoring arguments in json file if provided in command line
if args.json_args:
l = args.json_args
if "--cpu" in sys.argv[1:] or "--gpu" in sys.argv[1:]:
l = list(filter(lambda a: a not in ("--cpu", "--gpu"), l))
if "--auto-resize" in sys.argv[1:] or "--auto-resize-crop" in sys.argv[1:] \
or "--auto-rescale" in sys.argv[1:] or "--overlay" in sys.argv[1:]:
l = list(filter(lambda a: a not in ("--auto-resize",
"--auto-resize-crop", "--auto-rescale", "--overlay"), l))
args = parser.parse_args(l + sys.argv[1:])
config_args(parser, args)
args.func(args)
"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"
)
gpu_info_parser = subparsers.add_parser('gpu-info')
gpu_info_parser.add_argument(
"-j",
"--json",
default=False,
action="store_true",
help="Print GPU info as JSON"
)
# Register Command Handlers
ArgvParser.parser.set_defaults(func=main)
gpu_info_parser.set_defaults(func=gpu_info.main)
# 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)

5
config.py
Unescape View File

@ -1,6 +1,3 @@ @@ -1,6 +1,3 @@
import os
class Config:
"""
Variables Configuration Class
@ -55,4 +52,4 @@ class Config: @@ -55,4 +52,4 @@ class Config:
# Multiprocessing
@staticmethod
def multiprocessing():
return Config.args['gpu_ids'] is not None and Config.args['n_cores'] > 1
return Config.args['gpu_ids'] is None and Config.args['n_cores'] > 1

22
main.py
Unescape View File

@ -63,20 +63,20 @@ def select_phases(): @@ -63,20 +63,20 @@ def select_phases():
shift_step(shift_ending=1)
return phases
phases = [DressToCorrect(), CorrectToMask(), MaskToMaskref(),
MaskrefToMaskdet(), MaskdetToMaskfin(), MaskfinToNude()]
phases = [DressToCorrect, CorrectToMask, MaskToMaskref,
MaskrefToMaskdet, MaskdetToMaskfin, MaskfinToNude]
if conf.args['overlay']:
phases = add_tail(phases, ImageToResized())
phases = add_tail(phases, ImageToCrop())
phases = add_head(phases, ImageToOverlay())
phases = add_tail(phases, ImageToResized)
phases = add_tail(phases, ImageToCrop)
phases = add_head(phases, ImageToOverlay)
elif conf.args['auto_resize']:
phases = add_tail(phases, ImageToResized())
phases = add_tail(phases, ImageToResized)
elif conf.args['auto_resize_crop']:
phases = add_tail(phases, ImageToResizedCrop())
phases = add_tail(phases, ImageToResizedCrop)
elif conf.args['auto_rescale']:
phases = add_tail(phases, ImageToRescale())
elif not conf.args['folder']:
phases = add_tail(phases, ImageToRescale)
elif not os.path.isfile(conf.args['input']):
check_shape(read_image(conf.args['input']))
return phases
@ -87,7 +87,7 @@ def select_processing(): @@ -87,7 +87,7 @@ def select_processing():
:return:
"""
phases = select_phases()
if conf.args['folder']:
if os.path.isdir(conf.args['input']):
process = processing_image_folder(phases)
elif conf.args['n_runs'] != 1:
process = multiple_image_processing(phases, conf.args['n_runs'])
@ -133,4 +133,4 @@ def processing_image_folder(phases): @@ -133,4 +133,4 @@ def processing_image_folder(phases):
if __name__ == "__main__":
freeze_support()
# start_rook()
argv.run()
argv.ArgvParser.run()

203
processing/__init__.py
Unescape View File

@ -1,84 +1,131 @@ @@ -1,84 +1,131 @@
import json
import os
import pathlib
import shutil
import sys
import tempfile
import time
from json import JSONDecodeError
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
from utils import camel_case_to_str, cv2_supported_extension, read_image, write_image, json_to_argv
class Process:
"""
Abstract Process Class
"""
def __init__(self):
def __init__(self, *_args, args=None):
"""
Process Constructor
:param args: <dict> args parameter to run the image transformation (default use conf.args)
"""
self.__start = time.time()
self._args = conf.args.copy() if args is None else args.copy()
def run(self):
self.info_start_run()
self.setup()
self.execute()
self.clean()
self.info_end_run()
"""
Run the process
:return: None
"""
self._info_start_run()
self._setup()
self._execute()
self._clean()
self._info_end_run()
def info_start_run(self):
def _info_start_run(self):
"""
Logging when the process run begin
:return: None
"""
self.__start = time.time()
conf.log.info("Executing {}".format(camel_case_to_str(self.__class__.__name__)))
def info_end_run(self):
def _info_end_run(self):
"""
Logging when the process run end
:return: None
"""
conf.log.info("{} Finish".format(camel_case_to_str(self.__class__.__name__)))
conf.log.debug("{} Done in {} seconds".format(
camel_case_to_str(self.__class__.__name__), round(time.time() - self.__start, 2)))
def setup(self):
def _setup(self):
"""
Setup the process to be ready to execute
:return: None
"""
pass
def execute(self):
def _execute(self):
"""
Execute the process
:return: None
"""
pass
def clean(self):
pass
def _clean(self):
"""
Cleanup a process execution
:return: None
"""
def __str__(self):
return str(self.__class__.__name__)
class SimpleTransform(Process):
def __init__(self, *args):
super().__init__()
"""
Simple Transform Class
"""
def __new__(cls, input_path, phases, output_path):
def __init__(self, input_path, phases, output_path, args):
super().__init__(input_path, phases, output_path, args)
def __new__(cls, input_path, phases, output_path, args=None):
"""
Create the correct SimpleTransform object (ImageTransform or GiftTransform) corresponding to the input_path format
:param input_path: <string> original image path to process
:param output_path: <string> image path to write the result.
:param phases: <ImageTransform[]> list of Class transformation each image
:param args: <dict> args parameter to run the image transformation (default use conf.args)
:return: <ImageTransform|GiftTransform|None> SimpleTransform object corresponding to the input_path format
"""
if os.path.splitext(input_path)[1] == ".gif":
return GifTransform(input_path, phases, output_path)
return GifTransform(input_path, phases, output_path, args=args)
elif os.path.splitext(input_path)[1] in cv2_supported_extension():
return SimpleImageTransform(input_path, phases, output_path)
return ImageTransform(input_path, phases, output_path, args=args)
else:
return None
class SimpleImageTransform(Process):
class ImageTransform(Process):
"""
Simple Image Processing Class
Image Processing Class
"""
def __init__(self, input_path, phases, output_path):
def __init__(self, input_path, phases, output_path, args=None):
"""
ProcessImage Constructor
:param input_path: <string> original image path to process
:param output_path: <string> image path to write the result.
:param phases: <ImageTransform[]> list of transformation each image
:param args: <dict> args parameter to run the image transformation (default use conf.args)
:param phases: <ImageTransform[]> list Class of transformation each image
"""
super().__init__()
super().__init__(args=args)
self.__phases = phases
self.__output_path = output_path
self.__altered_path = conf.args['altered']
self.__starting_step = conf.args['steps'][0] if conf.args['steps'] else 0
self.__ending_step = conf.args['steps'][1] if conf.args['steps'] else None
self.__altered_path = self._args['altered']
self.__starting_step = self._args['steps'][0] if self._args['steps'] else 0
self.__ending_step = self._args['steps'][1] if self._args['steps'] else None
conf.log.debug("All Phases : {}".format(self.__phases))
conf.log.debug("To Be Executed Phases : {}".format(self.__phases[self.__starting_step:self.__ending_step]))
@ -88,11 +135,11 @@ class SimpleImageTransform(Process): @@ -88,11 +135,11 @@ class SimpleImageTransform(Process):
for p in self.__phases[:self.__starting_step]
]
def info_start_run(self):
super().info_start_run()
def _info_start_run(self):
super()._info_start_run()
conf.log.info("Processing on {}".format(str(self.__image_steps)[2:-2]))
def setup(self):
def _setup(self):
try:
self.__image_steps = [read_image(x) if isinstance(x, str) else x for x in self.__image_steps]
except FileNotFoundError as e:
@ -103,22 +150,22 @@ class SimpleImageTransform(Process): @@ -103,22 +150,22 @@ class SimpleImageTransform(Process):
"directory path that contains valid images.")
sys.exit(1)
def execute(self):
def _execute(self):
"""
Execute all phases on the image
:return: None
"""
for p in self.__phases[len(self.__image_steps) - 1:]:
for p in (x(args=self._args) for x in self.__phases[len(self.__image_steps) - 1:]):
r = p.run(*[self.__image_steps[i] for i in p.input_index])
self.__image_steps.append(r)
if self.__altered_path:
write_image(r, os.path.join(self.__altered_path, "{}.png".format(p.__class__.__name__)))
conf.log.debug("Writing {}, Result of the Execution of {}"
.format(
os.path.join(self.__altered_path, "{}.png".format(p.__class__.__name__)),
camel_case_to_str(p.__class__.__name__),
))
conf.log.debug("{} Step Image Of {} Execution"
.format(
os.path.join(self.__altered_path, "{}.png".format(p.__class__.__name__)),
camel_case_to_str(p.__class__.__name__),
))
write_image(self.__image_steps[-1], self.__output_path)
conf.log.info("{} Created".format(self.__output_path))
@ -133,15 +180,15 @@ class MultipleImageTransform(Process): @@ -133,15 +180,15 @@ class MultipleImageTransform(Process):
Multiple Image Processing Class
"""
def __init__(self, input_paths, phases, output_paths, children_process=SimpleTransform):
def __init__(self, input_paths, phases, output_paths, children_process=SimpleTransform, args=None):
"""
ProcessMultipleImages Constructor
:param input_paths: <string[]> images path list to process
:param output_paths: <string> images path to write the result
:param children_process: <ImageTransform> Process to use on the list of input
:param phases: <ImageTransform[]> list of transformation use by the process each image
:param phases: <ImageTransform[]> list of Class transformation use by the process each image
"""
super().__init__()
super().__init__(args=args)
self._phases = phases
self._input_paths = input_paths
self._output_paths = output_paths
@ -149,19 +196,18 @@ class MultipleImageTransform(Process): @@ -149,19 +196,18 @@ class MultipleImageTransform(Process):
self.__multiprocessing = conf.multiprocessing()
self.__children_process = children_process
def setup(self):
# TODO detect GIF or JPEG
self._process_list = [self.__children_process(i[0], self._phases, i[1])
def _setup(self):
self._process_list = [self.__children_process(i[0], self._phases, i[1], args=self._args)
for i in zip(self._input_paths, self._output_paths)]
def execute(self):
def _execute(self):
"""
Execute all phases on the list of images
:return: None
"""
def process_one_image(a):
conf.log.info("Processing image : {}/{}".format(a[1] + 1, len(self._process_list)))
conf.log.info("Processing Image : {}/{}".format(a[1] + 1, len(self._process_list)))
a[0].run()
if not self.__multiprocessing:
@ -180,47 +226,80 @@ class FolderImageTransform(MultipleImageTransform): @@ -180,47 +226,80 @@ class FolderImageTransform(MultipleImageTransform):
Folder Image Processing Class
"""
def __init__(self, input_folder_path, phases, output_folder_path):
def __init__(self, input_folder_path, phases, output_folder_path, args=None):
"""
FolderImageTransform Constructor
"""
super().__init__([], phases, [])
super().__init__([], phases, [], args=args)
self.__input_folder_path = input_folder_path
self.__output_folder_path = output_folder_path
self.__multiprocessing = conf.multiprocessing()
def setup(self):
def _setup(self):
conf.log.debug([(r, d, f) for r, d, f in os.walk(self.__input_folder_path)])
self._process_list = [
MultipleImageTransform(
[
x.path for x in os.scandir(os.path.join(r))
x.path for x in os.scandir(r)
if x.is_file() and os.path.splitext(x.path)[1] in cv2_supported_extension() + [".gif"]
],
self._phases,
[
"{}{}{}".format(os.path.splitext(x.path)[0], '_out', os.path.splitext(x.path)[1])
if not conf.args['output'] else os.path.join(conf.args['output'], r, os.path.basename(x.path))
for x in os.scandir(os.path.join(r))
"{}{}{}".format(
os.path.splitext(x.path)[0],
'_out',
os.path.splitext(x.path)[1]
)
if not conf.args['output'] else
os.path.join(
conf.args['output'],
pathlib.Path(*pathlib.Path(r).parts[1:]),
os.path.basename(x.path)
)
for x in os.scandir(r)
if x.is_file() and os.path.splitext(x.path)[1] in cv2_supported_extension() + [".gif"]
]
],
args=self.__get_folder_args(r)
) for r, _, _ in os.walk(self.__input_folder_path)
]
def __get_folder_args(self, folder_path):
json_path = os.path.join(folder_path, self._args['json_folder_name'])
conf.log.debug("Json Path Setting Path: {}".format(json_path))
if not os.path.isfile(json_path):
conf.log.info("No Json File Settings Found In {}. Using Default Configuration. ".format(folder_path))
return self._args
try:
with open(json_path, 'r') as f:
json_data = json.load(f)
except JSONDecodeError:
conf.log.info("Json File Settings {} Is Not In Valid JSON Format. Using Default Configuration. "
.format(folder_path))
return self._args
try:
a = argv.ArgvParser.config_args(argv.ArgvParser.parser.parse_args(sys.argv[1:]), json_data=json_data)
conf.log.info("Using {} Configuration for processing {} folder. "
.format(json_path, folder_path))
return a
except SystemExit:
conf.log.error("Arguments json file {} contains configuration error. "
"Using Default Configuration".format(json_path))
return self._args
class GifTransform(Process):
"""
GIF Image Processing Class
"""
def __init__(self, input_path, phases, output_path):
def __init__(self, input_path, phases, output_path, args=None):
"""
ImageTransformGIF Constructor
:param images: <string> gif path to process
:param input_path: <string> gif path to process
:param output_path: <string> image path to write the result
:param phases: <ImageTransform[]> list of transformation use by the process each image
:param phases: <ImageTransform[]> list of Class transformation use by the process each image
"""
super().__init__()
super().__init__(args=args)
self.__phases = phases
self.__input_path = input_path
self.__output_path = output_path
@ -228,7 +307,7 @@ class GifTransform(Process): @@ -228,7 +307,7 @@ class GifTransform(Process):
self.__temp_input_paths = []
self.__temp_output_paths = []
def setup(self):
def _setup(self):
self.__tmp_dir = tempfile.mkdtemp()
conf.log.debug("Temporay dir is {}".format(self.__tmp_dir))
imgs = imageio.mimread(self.__input_path)
@ -241,15 +320,19 @@ class GifTransform(Process): @@ -241,15 +320,19 @@ class GifTransform(Process):
[write_image(cv2.cvtColor(i[0], cv2.COLOR_RGB2BGR), i[1]) for i in zip(imgs, self.__temp_input_paths)]
def execute(self):
def _execute(self):
"""
Execute all phases on each frames of the gif and recreate the gif
:return: None
"""
MultipleImageTransform(self.__temp_input_paths, self.__phases, self.__temp_output_paths).run()
MultipleImageTransform(self.__temp_input_paths, self.__phases, self.__temp_output_paths, args=self._args).run()
dir = os.path.dirname(self.__output_path)
if dir != '':
os.makedirs(os.path.dirname(self.__output_path), exist_ok=True)
imageio.mimsave(self.__output_path, [imageio.imread(i) for i in self.__temp_output_paths])
conf.log.info("{} Gif Created ".format(self.__output_path))
def clean(self):
def _clean(self):
shutil.rmtree(self.__tmp_dir)

2
requirements.txt
Unescape View File

@ -6,6 +6,4 @@ torchvision==0.2.2.post3 @@ -6,6 +6,4 @@ torchvision==0.2.2.post3
torch==1.1.0
imageio==2.5.0
python-dotenv==0.10.3
rook==0.1.73
sentry-sdk==0.10.2
coloredlogs==10.0

6
transform/__init__.py
Unescape View File

@ -12,13 +12,15 @@ class ImageTransform: @@ -12,13 +12,15 @@ class ImageTransform:
Abstract Image Transformation Class
"""
def __init__(self, input_index=(-1,)):
def __init__(self, input_index=(-1,), args=None):
"""
Image Transformation Class Constructor
:param input_index: index where to take the input (default is -1 for previous transformation)
:param input_index: <tuple> index where to take the inputs (default is (-1) for previous transformation)
:param args: <dict> args parameter to run the image transformation (default use conf.args)
"""
self.__start = time.time()
self.input_index = input_index
self._args = conf.args.copy() if args is None else args.copy()
def run(self, *args):
self.__start = time.time()

8
transform/gan/__init__.py
Unescape View File

@ -18,16 +18,18 @@ class ImageTransformGAN(ImageTransform): @@ -18,16 +18,18 @@ class ImageTransformGAN(ImageTransform):
Abstract GAN Image Transformation Class
"""
def __init__(self, checkpoint, phase):
def __init__(self, checkpoint, phase, input_index=(-1,), args=None):
"""
Abstract GAN Image Transformation Class Constructor
:param checkpoint: <string> path to the checkpoint
:param phase: <string> phase name
:param input_index: <tuple> index where to take the inputs (default is (-1) for previous transformation)
:param args: <dict> args parameter to run the image transformation (default use conf.args)
"""
super().__init__()
super().__init__(input_index=input_index, args=args)
self.__checkpoint = checkpoint
self.__phase = phase
self.__gpu_ids = conf.args["gpu_ids"]
self.__gpu_ids = self._args["gpu_ids"]
def setup(self, image):
"""

33
transform/gan/mask.py
Unescape View File

@ -7,11 +7,18 @@ class CorrectToMask(ImageTransformGAN): @@ -7,11 +7,18 @@ class CorrectToMask(ImageTransformGAN):
Correct -> Mask [GAN]
"""
def __init__(self):
def __init__(self, input_index=(-1,), args=None):
"""
CorrectToMask Constructor
:param input_index: <tuple> index where to take the inputs (default is (-1) for previous transformation)
:param args: <dict> args parameter to run the image transformation (default use conf.args)
"""
super().__init__(conf.args['checkpoints']["correct_to_mask"], "correct_to_mask")
super().__init__(
(args if args is not None else conf.args)['checkpoints']["correct_to_mask"],
"correct_to_mask",
input_index=input_index,
args=args
)
class MaskrefToMaskdet(ImageTransformGAN):
@ -19,11 +26,18 @@ class MaskrefToMaskdet(ImageTransformGAN): @@ -19,11 +26,18 @@ class MaskrefToMaskdet(ImageTransformGAN):
Maskref -> Maskdet [GAN]
"""
def __init__(self):
def __init__(self, input_index=(-1,), args=None):
"""
MaskrefToMaskdet Constructor
:param input_index: <tuple> index where to take the inputs (default is (-1) for previous transformation)
:param args: <dict> args parameter to run the image transformation (default use conf.args)
"""
super().__init__(conf.args['checkpoints']["maskref_to_maskdet"], "maskref_to_maskdet")
super().__init__(
(args if args is not None else conf.args)['checkpoints']["maskref_to_maskdet"],
"maskref_to_maskdet",
input_index=input_index,
args=args
)
class MaskfinToNude(ImageTransformGAN):
@ -31,8 +45,15 @@ class MaskfinToNude(ImageTransformGAN): @@ -31,8 +45,15 @@ class MaskfinToNude(ImageTransformGAN):
Maskfin -> Nude [GAN]
"""
def __init__(self):
def __init__(self, input_index=(-1,), args=None):
"""
MaskfinToNude Constructor
:param input_index: <tuple> index where to take the inputs (default is (-1) for previous transformation)
:param args: <dict> args parameter to run the image transformation (default use conf.args)
"""
super().__init__(conf.args['checkpoints']["maskfin_to_nude"], "maskfin_to_nude")
super().__init__(
(args if args is not None else conf.args)['checkpoints']["maskfin_to_nude"],
"maskfin_to_nude",
input_index=input_index,
args=args
)

9
transform/opencv/__init__.py
Unescape View File

@ -6,8 +6,13 @@ class ImageTransformOpenCV(ImageTransform): @@ -6,8 +6,13 @@ class ImageTransformOpenCV(ImageTransform):
OPENCV Image Transform class
"""
def __init__(self, input_index=(-1,)):
super().__init__(input_index)
def __init__(self, input_index=(-1,), args=None,):
"""
ImageTransformOpenCV Constructor
:param input_index: <tuple> index where to take the inputs (default is (-1) for previous transformation)
:param args: <dict> args parameter to run the image transformation (default use conf.args)
"""
super().__init__(args=args, input_index=input_index)
class BodyPart:

22
transform/opencv/mask.py
Unescape View File

@ -10,10 +10,12 @@ from config import Config as conf @@ -10,10 +10,12 @@ from config import Config as conf
class MaskToMaskref(ImageTransformOpenCV):
"""
Mask & Correct -> MaskRef [OPENCV]
:param input_index: <tuple> index where to take the inputs (default is (-2,-1) for the two previous transformation)
:param args: <dict> args parameter to run the image transformation (default use conf.args)
"""
def __init__(self, input_index=(-2, -1)):
super().__init__(input_index=input_index)
def __init__(self, input_index=(-2, -1), args=None):
super().__init__(args=args, input_index=input_index)
def execute(self, correct, mask):
"""
@ -52,15 +54,17 @@ class MaskToMaskref(ImageTransformOpenCV): @@ -52,15 +54,17 @@ class MaskToMaskref(ImageTransformOpenCV):
class MaskdetToMaskfin(ImageTransformOpenCV):
"""
Maskdet -> Maskfin [OPENCV]
:param input_index: <tuple> index where to take the inputs (default is (-2,-1) for the two previous transformation)
:param args: <dict> args parameter to run the image transformation (default use conf.args)
"""
def __init__(self, input_index=(-2, -1)):
super().__init__(input_index=input_index)
self.__aur_size = conf.args["prefs"]["aursize"]
self.__nip_size = conf.args["prefs"]["nipsize"]
self.__tit_size = conf.args["prefs"]["titsize"]
self.__vag_size = conf.args["prefs"]["vagsize"]
self.__hair_size = conf.args["prefs"]["hairsize"]
def __init__(self, input_index=(-2, -1), args=None,):
super().__init__(input_index=input_index, args=args)
self.__aur_size = self._args["prefs"]["aursize"]
self.__nip_size = self._args["prefs"]["nipsize"]
self.__tit_size = self._args["prefs"]["titsize"]
self.__vag_size = self._args["prefs"]["vagsize"]
self.__hair_size = self._args["prefs"]["hairsize"]
def execute(self, maskref, maskdet):
"""

28
transform/opencv/resize.py
Unescape View File

@ -9,13 +9,15 @@ from transform.opencv.correct import DressToCorrect @@ -9,13 +9,15 @@ from transform.opencv.correct import DressToCorrect
class ImageToCrop(ImageTransformOpenCV):
"""
Image -> Crop [OPENCV]
:param input_index: <tuple> index where to take the inputs (default is (-1) for previous transformation)
:param args: <dict> args parameter to run the image transformation (default use conf.args)
"""
def __init__(self, input_index=(-1,)):
super().__init__(input_index=input_index)
self.__x1 = conf.args['overlay'][0]
self.__y1 = conf.args['overlay'][1]
self.__x2 = conf.args['overlay'][2]
self.__y2 = conf.args['overlay'][3]
def __init__(self, input_index=(-1,), args=None):
super().__init__(args=args, input_index=input_index)
self.__x1 = self._args['overlay'][0]
self.__y1 = self._args['overlay'][1]
self.__x2 = self._args['overlay'][2]
self.__y2 = self._args['overlay'][3]
def execute(self, img):
"""
@ -33,13 +35,15 @@ class ImageToCrop(ImageTransformOpenCV): @@ -33,13 +35,15 @@ class ImageToCrop(ImageTransformOpenCV):
class ImageToOverlay(ImageTransformOpenCV):
"""
Image -> Overlay [OPENCV]
:param input_index: <tuple> index where to take the inputs (default is (0,1) for first and previous transformation)
:param args: <dict> args parameter to run the image transformation (default use conf.args)
"""
def __init__(self, input_index=(0, -1)):
super().__init__(input_index=input_index)
self.__x1 = conf.args['overlay'][0]
self.__y1 = conf.args['overlay'][1]
self.__x2 = conf.args['overlay'][2]
self.__y2 = conf.args['overlay'][3]
def __init__(self, input_index=(0, -1), args=None):
super().__init__(input_index=input_index, args=args,)
self.__x1 = self._args['overlay'][0]
self.__y1 = self._args['overlay'][1]
self.__x2 = self._args['overlay'][2]
self.__y2 = self._args['overlay'][3]
def execute(self, img_to_overlay, img):
"""

10
transform/opencv/watermark.py
Unescape View File

@ -2,6 +2,7 @@ import cv2 @@ -2,6 +2,7 @@ import cv2
import numpy as np
from transform.opencv import ImageTransformOpenCV
from config import Config as conf
class ImageToWatermark(ImageTransformOpenCV):
@ -9,8 +10,13 @@ class ImageToWatermark(ImageTransformOpenCV): @@ -9,8 +10,13 @@ class ImageToWatermark(ImageTransformOpenCV):
Image -> Watermarked Image [OPENCV]
"""
def __init__(self, watermark="fake.png"):
super().__init__()
def __init__(self, input_index=(-1,), args=None, watermark="fake.png"):
"""
:param input_index: <tuple> index where to take the inputs (default is (-1) for previous transformation)
:param args: <dict> args parameter to run the image transformation (default use conf.args)
:param watermark: <string> path to the watermark image
"""
super().__init__(args=args, input_index=input_index)
self.__watermark = cv2.imread(watermark, cv2.IMREAD_UNCHANGED)
def execute(self, img):
</