diff --git a/argv.py b/argv.py
index eedd2d5..b461fec 100644
--- a/argv.py
+++ b/argv.py
@@ -9,6 +9,7 @@ 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):
@@ -45,20 +46,26 @@ def config_args(parser, args):
def config_args_in():
if not args.input:
parser.error("-i, --input INPUT is required.")
- elif not os.path.isfile(args.input):
- parser.error(" {} file doesn't exist".format(args.input))
+ 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.output:
+ 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))
+ parser.error("{} directory doesn't exist.".format(args.input))
if args.func == main:
conf.args = vars(args)
@@ -90,6 +97,13 @@ def run():
"--output",
help="path where the transformed photo will be saved. (default: output.)",
)
+ 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",
@@ -135,9 +149,6 @@ def run():
default=0,
help="Pubic hair size scalar best results set to 0 to disable",
)
- parser.add_argument(
- "--gif", action="store_true", default=False, help="run the processing on a gif"
- )
parser.add_argument(
"-n", "--n_runs", type=int, default=1, help="number of times to process input (default: 1)",
)
@@ -165,16 +176,6 @@ def run():
help="Scale image to 512x512",
)
- 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"
- )
-
def check_crops_coord():
def type_func(a):
if not re.match(r"^\d+,\d+:\d+,\d+$", a):
@@ -191,36 +192,6 @@ def run():
"(,:,) and overlay the result on the original image.",
)
- 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 = {}
- try:
- data = json.load(f)
- 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.",
- )
-
def check_steps_args():
def type_func(a):
if not re.match(r"^[0-5]:[0-5]$", a):
@@ -249,7 +220,6 @@ def run():
"--steps",
type=check_steps_args(),
help="Select a range of steps to execute :."
- "Scale options are ignored atm when using this option"
"Steps are : \n"
"0 : dress -> correct [OPENCV]\n"
"1 : correct -> mask [GAN]\n"
@@ -268,10 +238,48 @@ def run():
parser.add_argument(
"-c",
"--checkpoints",
- default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "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):
+ if not os.path.isfile(a):
+ raise argparse.ArgumentTypeError(
+ "Arguments json file {} not found.".format(a))
+ with open(a) as f:
+ data = {}
+ try:
+ data = json.load(f)
+ 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)
diff --git a/main.py b/main.py
index e5e126f..36ecb9e 100644
--- a/main.py
+++ b/main.py
@@ -8,8 +8,7 @@ import argv
from config import Config as conf
from utils import setup_log, read_image, check_shape
-from processing.gif import SimpleGIFTransform
-from processing.image import SimpleImageTransform, MultipleImageTransform
+from processing import SimpleTransform, FolderImageTransform, MultipleImageTransform
from transform.gan.mask import CorrectToMask, MaskrefToMaskdet, MaskfinToNude
from transform.opencv.resize import ImageToCrop, ImageToOverlay, ImageToRescale, ImageToResized, ImageToResizedCrop
from transform.opencv.correct import DressToCorrect
@@ -77,7 +76,7 @@ def select_phases():
phases = add_tail(phases, ImageToResizedCrop())
elif conf.args['auto_rescale']:
phases = add_tail(phases, ImageToRescale())
- else:
+ elif not conf.args['folder']:
check_shape(read_image(conf.args['input']))
return phases
@@ -88,10 +87,8 @@ def select_processing():
:return:
"""
phases = select_phases()
- if conf.args['gif'] and conf.args['n_runs'] != 1:
- process = multiple_gif_processing(phases, conf.args['n_runs'])
- elif conf.args['gif']:
- process = simple_gif_processing(phases)
+ if conf.args['folder']:
+ process = processing_image_folder(phases)
elif conf.args['n_runs'] != 1:
process = multiple_image_processing(phases, conf.args['n_runs'])
else:
@@ -100,38 +97,13 @@ def select_processing():
return process
-def simple_gif_processing(phases):
- """
- Define a simple gif process ready to run
- :param phases: list of image transformation
- :return: a gif process run ready
- """
- return SimpleGIFTransform(conf.args['input'], phases, conf.args['output'])
-
-
-def multiple_gif_processing(phases, n):
- """
- Define a multiple gif process ready to run
- :param phases: list of image transformation
- :param n: number of times to process
- :return: a multiple gif process run ready
- """
- filename, extension = os.path.splitext(conf.args['output'])
- return MultipleImageTransform(
- [conf.args['input'] for _ in range(n)],
- phases,
- ["{}{}{}".format(filename, i, extension) for i in range(n)],
- SimpleGIFTransform
- )
-
-
def simple_image_processing(phases):
"""
Define a simple image process ready to run
:param phases: list of image transformation
- :return: a image process run ready
+ :return: a image process run ready
"""
- return SimpleImageTransform(conf.args['input'], phases, conf.args['output'])
+ return SimpleTransform(conf.args['input'], phases, conf.args['output'])
def multiple_image_processing(phases, n):
@@ -149,6 +121,15 @@ def multiple_image_processing(phases, n):
)
+def processing_image_folder(phases):
+ """
+ Define a folder image process ready to run
+ :param phases: list of image transformation
+ :return: a image process run ready
+ """
+ return FolderImageTransform(conf.args['input'], phases, conf.args['output'])
+
+
if __name__ == "__main__":
freeze_support()
# start_rook()
diff --git a/processing/__init__.py b/processing/__init__.py
index 762d86f..4954255 100644
--- a/processing/__init__.py
+++ b/processing/__init__.py
@@ -1,14 +1,21 @@
+import os
+import shutil
+import sys
+import tempfile
import time
+from multiprocessing.pool import ThreadPool
+
+import cv2
+import imageio
from config import Config as conf
-from utils import camel_case_to_str
+from utils import camel_case_to_str, cv2_supported_extension, read_image, write_image
class Process:
"""
Abstract Process Class
"""
-
def __init__(self):
self.__start = time.time()
@@ -38,4 +45,211 @@ class Process:
pass
def __str__(self):
- return str(self.__class__.__name__)
\ No newline at end of file
+ return str(self.__class__.__name__)
+
+
+class SimpleTransform(Process):
+ def __init__(self, *args):
+ super().__init__()
+
+ def __new__(cls, input_path, phases, output_path):
+ if os.path.splitext(input_path)[1] == ".gif":
+ return GifTransform(input_path, phases, output_path)
+ elif os.path.splitext(input_path)[1] in cv2_supported_extension():
+ return SimpleImageTransform(input_path, phases, output_path)
+ else:
+ return None
+
+
+class SimpleImageTransform(Process):
+ """
+ Simple Image Processing Class
+ """
+
+ def __init__(self, input_path, phases, output_path):
+ """
+ ProcessImage Constructor
+ :param input_path: original image path to process
+ :param output_path: image path to write the result.
+ :param phases: list of transformation each image
+ """
+ super().__init__()
+ 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
+
+ conf.log.debug("All Phases : {}".format(self.__phases))
+ conf.log.debug("To Be Executed Phases : {}".format(self.__phases[self.__starting_step:self.__ending_step]))
+
+ self.__image_steps = [input_path] + [
+ os.path.join(self.__altered_path, "{}.png".format(p.__class__.__name__))
+ for p in self.__phases[:self.__starting_step]
+ ]
+
+ def info_start_run(self):
+ super().info_start_run()
+ conf.log.info("Processing on {}".format(str(self.__image_steps)[2:-2]))
+
+ 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:
+ conf.log.error(e)
+ conf.log.error("{} is not able to resume because it not able to load required images. "
+ .format(camel_case_to_str(self.__class__.__name__)))
+ conf.log.error("Possible source of this error is that --altered argument is not a correct "
+ "directory path that contains valid images.")
+ sys.exit(1)
+
+ def execute(self):
+ """
+ Execute all phases on the image
+ :return: None
+ """
+ for p 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__),
+ ))
+
+ write_image(self.__image_steps[-1], self.__output_path)
+ conf.log.info("{} Created".format(self.__output_path))
+ conf.log.debug("{} Result Image Of {} Execution"
+ .format(self.__output_path, camel_case_to_str(self.__class__.__name__)))
+
+ return self.__image_steps[-1]
+
+
+class MultipleImageTransform(Process):
+ """
+ Multiple Image Processing Class
+ """
+
+ def __init__(self, input_paths, phases, output_paths, children_process=SimpleTransform):
+ """
+ ProcessMultipleImages Constructor
+ :param input_paths: images path list to process
+ :param output_paths: images path to write the result
+ :param children_process: Process to use on the list of input
+ :param phases: list of transformation use by the process each image
+ """
+ super().__init__()
+ self._phases = phases
+ self._input_paths = input_paths
+ self._output_paths = output_paths
+ self._process_list = []
+ 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])
+ for i in zip(self._input_paths, self._output_paths)]
+
+ 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)))
+ a[0].run()
+
+ if not self.__multiprocessing:
+ for x in zip(self._process_list, range(len(self._process_list))):
+ process_one_image(x)
+ else:
+ conf.log.debug("Using Multiprocessing")
+ pool = ThreadPool(conf.args['n_cores'])
+ pool.map(process_one_image, zip(self._process_list, range(len(self._process_list))))
+ pool.close()
+ pool.join()
+
+
+class FolderImageTransform(MultipleImageTransform):
+ """
+ Folder Image Processing Class
+ """
+
+ def __init__(self, input_folder_path, phases, output_folder_path):
+ """
+ FolderImageTransform Constructor
+ """
+ super().__init__([], phases, [])
+ self.__input_folder_path = input_folder_path
+ self.__output_folder_path = output_folder_path
+ self.__multiprocessing = conf.multiprocessing()
+
+ 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))
+ 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))
+ if x.is_file() and os.path.splitext(x.path)[1] in cv2_supported_extension() + [".gif"]
+ ]
+ ) for r, _, _ in os.walk(self.__input_folder_path)
+ ]
+
+
+class GifTransform(Process):
+ """
+ GIF Image Processing Class
+ """
+
+ def __init__(self, input_path, phases, output_path):
+ """
+ ImageTransformGIF Constructor
+ :param images: gif path to process
+ :param output_path: image path to write the result
+ :param phases: list of transformation use by the process each image
+ """
+ super().__init__()
+ self.__phases = phases
+ self.__input_path = input_path
+ self.__output_path = output_path
+ self.__tmp_dir = None
+ self.__temp_input_paths = []
+ self.__temp_output_paths = []
+
+ def setup(self):
+ self.__tmp_dir = tempfile.mkdtemp()
+ conf.log.debug("Temporay dir is {}".format(self.__tmp_dir))
+ imgs = imageio.mimread(self.__input_path)
+ conf.log.info("GIF have {} Frames To Process".format(len(imgs)))
+ self.__temp_input_paths = [os.path.join(self.__tmp_dir, "intput_{}.png".format(i))
+ for i in range(len(imgs))]
+
+ self.__temp_output_paths = [os.path.join(self.__tmp_dir, "output_{}.png".format(i))
+ for i in range(len(imgs))]
+
+ [write_image(cv2.cvtColor(i[0], cv2.COLOR_RGB2BGR), i[1]) for i in zip(imgs, self.__temp_input_paths)]
+
+ 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()
+
+ 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):
+ shutil.rmtree(self.__tmp_dir)
diff --git a/processing/gif.py b/processing/gif.py
deleted file mode 100644
index fbcbbd7..0000000
--- a/processing/gif.py
+++ /dev/null
@@ -1,58 +0,0 @@
-import os
-import shutil
-import tempfile
-
-import cv2
-import imageio
-
-from config import Config as conf
-from processing import Process
-from processing.image import MultipleImageTransform
-from utils import write_image
-
-
-class SimpleGIFTransform(Process):
- """
- GIF Image Processing Class
- """
-
- def __init__(self, input_path, phases, output_path):
- """
- ImageTransformGIF Constructor
- :param images: gif path to process
- :param output_path: image path to write the result
- :param phases: list of transformation use by the process each image
- """
- super().__init__()
- self.__phases = phases
- self.__input_path = input_path
- self.__output_path = output_path
- self.__tmp_dir = None
- self.__temp_input_paths = []
- self.__temp_output_paths = []
-
- def setup(self):
- self.__tmp_dir = tempfile.mkdtemp()
- conf.log.debug("Temporay dir is {}".format(self.__tmp_dir))
- imgs = imageio.mimread(self.__input_path)
- conf.log.info("GIF have {} Frames To Process".format(len(imgs)))
- self.__temp_input_paths = [os.path.join(self.__tmp_dir, "intput_{}.png".format(i))
- for i in range(len(imgs))]
-
- self.__temp_output_paths = [os.path.join(self.__tmp_dir, "output_{}.png".format(i))
- for i in range(len(imgs))]
-
- [write_image(cv2.cvtColor(i[0], cv2.COLOR_RGB2BGR), i[1]) for i in zip(imgs, self.__temp_input_paths)]
-
- 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()
-
- 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):
- shutil.rmtree(self.__tmp_dir)
\ No newline at end of file
diff --git a/processing/image.py b/processing/image.py
deleted file mode 100644
index 8d53301..0000000
--- a/processing/image.py
+++ /dev/null
@@ -1,119 +0,0 @@
-import os
-import sys
-from multiprocessing.pool import ThreadPool
-
-from config import Config as conf
-from processing import Process
-from utils import read_image, write_image, camel_case_to_str
-
-
-class SimpleImageTransform(Process):
- """
- Simple Image Processing Class
- """
-
- def __init__(self, input_path, phases, output_path):
- """
- ProcessImage Constructor
- :param input_path: original image path to process
- :param output_path: image path to write the result.
- :param phases: list of transformation each image
- """
- super().__init__()
- 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
-
- conf.log.debug("All Phases : {}".format(self.__phases))
- conf.log.debug("To Be Executed Phases : {}".format(self.__phases[self.__starting_step:self.__ending_step]))
-
- self.__image_steps = [input_path] + [
- os.path.join(self.__altered_path, "{}.png".format(p.__class__.__name__))
- for p in self.__phases[:self.__starting_step]
- ]
-
- def info_start_run(self):
- super().info_start_run()
- conf.log.debug("Processing on {}".format(self.__image_steps))
-
- 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:
- conf.log.error(e)
- conf.log.error("{} is not able to resume because it not able to load required images. "
- .format(camel_case_to_str(self.__class__.__name__)))
- conf.log.error("Possible source of this error is that --altered argument is not a correct "
- "directory path that contains valid images.")
- sys.exit(1)
-
- def execute(self):
- """
- Execute all phases on the image
- :return: None
- """
- for p 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__),
- ))
-
- write_image(self.__image_steps[-1], self.__output_path)
- conf.log.debug("Writing {}, Result Image Of {} Execution"
- .format(self.__output_path, camel_case_to_str(self.__class__.__name__)))
-
- return self.__image_steps[-1]
-
-
-class MultipleImageTransform(Process):
- """
- Multiple Image Processing Class
- """
-
- def __init__(self, input_paths, phases, output_paths, children_process=SimpleImageTransform):
- """
- ProcessMultipleImages Constructor
- :param input_paths: images path list to process
- :param output_paths: images path to write the result
- :param children_process: Process to use on the list of input
- :param phases: list of transformation use by the process each image
- """
- super().__init__()
- self.__phases = phases
- self.__input_paths = input_paths
- self.__output_paths = output_paths
- self.__process_list = []
- self.__multiprocessing = conf.multiprocessing()
- self.__children_process = children_process
-
- def setup(self):
- self.__process_list = [self.__children_process(i[0], self.__phases, i[1])
- for i in zip(self.__input_paths, self.__output_paths)]
-
- 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)))
- a[0].run()
-
- if not self.__multiprocessing:
- for x in zip(self.__process_list, range(len(self.__process_list))):
- process_one_image(x)
- else:
- conf.log.debug("Using Multiprocessing")
- pool = ThreadPool(conf.args['n_cores'])
- pool.map(process_one_image, zip(self.__process_list, range(len(self.__process_list))))
- pool.close()
- pool.join()
diff --git a/utils.py b/utils.py
index b366aaf..eba8e99 100644
--- a/utils.py
+++ b/utils.py
@@ -40,6 +40,10 @@ def write_image(image, path):
if dir != '':
os.makedirs(os.path.dirname(path), exist_ok=True)
+ if os.path.splitext(path)[1] not in cv2_supported_extension():
+ conf.log.error("{} invalid extension format.".format(path))
+ sys.exit(1)
+
cv2.imwrite(path, image)
if not check_image_file_validity(path):
@@ -105,3 +109,8 @@ def camel_case_to_str(identifier):
"""
matches = finditer('.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)', identifier)
return " ".join([m.group(0) for m in matches])
+
+
+def cv2_supported_extension():
+ return [".bmp", ".dib", ".jpeg", ".jpg", ".jpe", ".jp2", ".png",
+ ".pbm", ".pgm", "ppm", ".sr", ".ras", ".tiff", ".tif"]