Browse Source

Initial folder processing support #30

* Add -f, --folder option for processing folder
* No support GIF yet
* No support for json file per folder yet
master
PommeDroid 3 years ago
parent
commit
a02547b2ca
  1. 106
      argv.py
  2. 15
      main.py
  3. 59
      processing/image.py
  4. 9
      utils.py

106
argv.py

@ -9,6 +9,7 @@ import gpu_info @@ -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,28 @@ def config_args(parser, args): @@ -45,20 +46,28 @@ 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 not args.gif and os.path.splitext(args.input)[1] not in cv2_supported_extension():
parser.error("Input {} file not supported format.".format(args.input))
elif not args.folder and args.gif and os.path.splitext(args.input)[1] != ".gif":
parser.error("Input {} file if not a gif.".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 +99,14 @@ def run(): @@ -90,6 +99,14 @@ def run():
"--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
"Gif not supported atm", # TODO Gif Support
)
processing_mod = parser.add_mutually_exclusive_group()
processing_mod.add_argument(
"--cpu",
@ -165,16 +182,6 @@ def run(): @@ -165,16 +182,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 +198,6 @@ def run(): @@ -191,36 +198,6 @@ def run():
"(<x_top_left>,<y_top_left>:<x_bot_right>,<x_bot_right>) 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 +226,6 @@ def run(): @@ -249,7 +226,6 @@ def run():
"--steps",
type=check_steps_args(),
help="Select a range of steps to execute <starting step>:<ending step>."
"Scale options are ignored atm when using this option"
"Steps are : \n"
"0 : dress -> correct [OPENCV]\n"
"1 : correct -> mask [GAN]\n"
@ -268,10 +244,48 @@ def run(): @@ -268,10 +244,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)

15
main.py

@ -9,7 +9,7 @@ from config import Config as conf @@ -9,7 +9,7 @@ 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.image import SimpleImageTransform, MultipleImageTransform, FolderImageTransform
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 +77,7 @@ def select_phases(): @@ -77,7 +77,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,7 +88,9 @@ def select_processing(): @@ -88,7 +88,9 @@ def select_processing():
:return:
"""
phases = select_phases()
if conf.args['gif'] and conf.args['n_runs'] != 1:
if conf.args['folder']:
process = processing_image_folder(phases)
elif 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)
@ -148,6 +150,13 @@ def multiple_image_processing(phases, n): @@ -148,6 +150,13 @@ def multiple_image_processing(phases, n):
["{}{}{}".format(filename, i, extension) for i in range(n)]
)
def processing_image_folder(phases):
"""
Define a folder image process ready to run
:param phases: <ImageTransform[]> list of image transformation
:return: <SimpleImageTransform> a image process run ready
"""
return FolderImageTransform(conf.args['input'], phases, conf.args['output'])
if __name__ == "__main__":
freeze_support()

59
processing/image.py

@ -4,7 +4,7 @@ from multiprocessing.pool import ThreadPool @@ -4,7 +4,7 @@ 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
from utils import read_image, write_image, camel_case_to_str, cv2_supported_extension
class SimpleImageTransform(Process):
@ -36,7 +36,7 @@ class SimpleImageTransform(Process): @@ -36,7 +36,7 @@ class SimpleImageTransform(Process):
def info_start_run(self):
super().info_start_run()
conf.log.debug("Processing on {}".format(self.__image_steps))
conf.log.info("Processing on {}".format(str(self.__image_steps)[2:-2]))
def setup(self):
try:
@ -44,7 +44,7 @@ class SimpleImageTransform(Process): @@ -44,7 +44,7 @@ class SimpleImageTransform(Process):
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__)))
.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)
@ -67,7 +67,8 @@ class SimpleImageTransform(Process): @@ -67,7 +67,8 @@ class SimpleImageTransform(Process):
))
write_image(self.__image_steps[-1], self.__output_path)
conf.log.debug("Writing {}, Result Image Of {} Execution"
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]
@ -87,16 +88,16 @@ class MultipleImageTransform(Process): @@ -87,16 +88,16 @@ class MultipleImageTransform(Process):
:param phases: <ImageTransform[]> 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._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)]
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):
"""
@ -105,15 +106,47 @@ class MultipleImageTransform(Process): @@ -105,15 +106,47 @@ class MultipleImageTransform(Process):
"""
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:
for x in zip(self.__process_list, range(len(self.__process_list))):
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.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()],
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()
]
) for r, _, _ in os.walk(self.__input_folder_path)
]

9
utils.py

@ -40,6 +40,10 @@ def write_image(image, path): @@ -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): @@ -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"

Loading…
Cancel
Save