Browse Source

Heavy refactoring

* Change all the logic of the phases to more flexiblility
* Add rescale/resize/overlay as phases (the same level that gan/opencv
masks transformation)
* Add more comments everywhere
* Modify Configuration(Options) Class to be more general (args, checkpoints
location and multingthread configuration are now acccesble from it)
* Add logging with debug option (-d, --debug)
* No fonctionnilaty change
tags/v1.2.10
PommeDroid 2 years ago
parent
commit
ec98f8666f

+ 17
- 0
.gitignore View File

@@ -5,6 +5,7 @@ logs
# Process Images
*.png
*.jpg
*.gif

# Runtime data
pids
@@ -32,3 +33,19 @@ dist/

# Output file
output.png

# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class

# Environments
.venv
env/
venv/
ENV/
env.bak/
venv.bak/

# IDE
.idea

+ 220
- 0
argv.py View File

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

import gpu_info
from main import main
from config import Config as conf


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

def set_gpu_ids():
gpu_ids = args.gpu
if args.cpu:
gpu_ids = None
elif gpu_ids is None:
gpu_ids = [0]
return gpu_ids

conf.args = vars(args)
conf.args['gpu_ids'] = set_gpu_ids()
conf.args['prefs'] = set_body_parts_prefs()


def check_args(args):
def check_args_in():
if not os.path.isfile(args.input):
print("Error : {} file doesn't exist".format(
args.input), file=sys.stderr)
exit(1)

check_args_in()


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", default="input.png", help="path of the photo to transform"
)
parser.add_argument(
"-o",
"--output",
default="output.png",
help="path where the transformed photo will be saved. (default: output.png or output.gif)",
)
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(
"--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)",
)
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",
)

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):
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('\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_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.",
)

# Register Command Handlers
parser.set_defaults(func=main)
gpu_info_parser.set_defaults(func=gpu_info.main)

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:])

check_args(args)
set_config_args(args)
args.func(args)

+ 65
- 0
config.py View File

@@ -0,0 +1,65 @@
import os


class Config:
"""
Variables Configuration Class
"""
# experiment specifics
norm = "batch" # instance normalization or batch normalization
use_dropout = False # use dropout for the generator
data_type = 32 # Supported data type i.e. 8, 16, 32 bit

# input/output sizes
batchSize = 1 # input batch size
input_nc = 3 # of input image channels
output_nc = 3 # of output image channels

# for setting inputs
# if true, takes images in order to make batches, otherwise takes them randomly
serial_batches = True
nThreads = (
0
) # threads for loading data. Keep this value at 0! see: https://github.com/pytorch/pytorch/issues/12831
# Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size,
# only a subset is loaded.
max_dataset_size = 1

# for generator
netG = "global" # selects model to use for netG
ngf = 64 # of gen filters in first conv layer
n_downsample_global = 4 # number of downsampling layers in netG
n_blocks_global = (
9
) # number of residual blocks in the global generator network
n_blocks_local = (
0
) # number of residual blocks in the local enhancer network
n_local_enhancers = 0 # number of local enhancers to use
# number of epochs that we only train the outmost local enhancer
niter_fix_global = 0

# Phase specific options
checkpoints_dir = ""
dataroot = ""

# Image requirement
desired_size = 512

# GAN checkpoints location
checkpoints = dict({
'correct_to_mask': os.path.join(os.path.dirname(os.path.realpath(__file__)), "checkpoints", "cm.lib"),
'maskref_to_maskdet': os.path.join(os.path.dirname(os.path.realpath(__file__)), "checkpoints", "mm.lib"),
'maskfin_to_nude': os.path.join(os.path.dirname(os.path.realpath(__file__)), "checkpoints", "mn.lib"),
})

# Argparser dict
args = {}

# Log
log = None

# Multiprocessing
@staticmethod
def multiprocessing():
return Config.args['gpu_ids'] is not None and Config.args['n_cores'] > 1

+ 108
- 324
main.py View File

@@ -1,346 +1,130 @@
import json
import re
import shutil
import logging
import os
import sys
import argparse
import tempfile
from json import JSONDecodeError

import cv2
import time
import os
import imageio

import sentry_sdk
import rook

import utils
import numpy as np

from run import process, process_gif
from multiprocessing import freeze_support
from multiprocessing.pool import ThreadPool
#from dotenv import load_dotenv

import gpu_info

#
#load_dotenv()

parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers()

parser.add_argument(
"-i", "--input", default="input.png", help="path of the photo to transform"
)
parser.add_argument(
"-o",
"--output",
default="output.png",
help="path where the transformed photo will be saved. (default: output.png or output.gif)",
)
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(
"--gif", action="store_true", default=False, help="run the processing on a gif"
)
parser.add_argument(
"-n", "--n_runs", type=int, help="number of times to process input (default: 1)",
)
parser.add_argument(
"--n_cores", type=int, default=4, help="number of cpu cores to use (default: 4)",
)

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",
)

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"
)

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

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('\d+', a))
from processing.gif import SimpleGIFTransform
from processing.image import SimpleImageTransform, MultipleImageTransform
from transform.gan.mask import CorrectToMask, MaskrefToMaskdet, MaskfinToMaskdet
from transform.opencv.resize import ImageToCrop, ImageToOverlay, ImageToRescale, ImageToResized, ImageToResizedCrop
from transform.opencv.correct import DressToCorrect
from transform.opencv.mask import MaskToMaskref, MaskdetToMaskfin

return type_func

def main(_):
"""
Main logic entry point
"""
conf.log = setup_log(logging.DEBUG) if conf.args['debug'] else setup_log()
conf.log.info("Welcome to DreamPower")

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_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.",
)

"""
main.py

How to run:
python3 main.py

"""

# ------------------------------------------------- main()


def main(args):
if not os.path.isfile(args.input):
print("Error : {} file doesn't exist".format(
args.input), file=sys.stderr)
exit(1)
start = time.time()

gpu_ids = args.gpu

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

if args.cpu:
gpu_ids = None
elif gpu_ids is None:
gpu_ids = [0]

if not args.gif:
# Read image
file = open(args.input, "rb")
image_bytes = bytearray(file.read())
np_image = np.asarray(image_bytes, dtype=np.uint8)
image = cv2.imdecode(np_image, cv2.IMREAD_COLOR)

# See if image loaded correctly
if image is None:
print("Error : {} file is not valid".format(
args.input), file=sys.stderr)
exit(1)

# Preprocess
if args.overlay:
original_image = image.copy()
image = utils.crop_input(
image, args.overlay[0], args.overlay[1], args.overlay[2], args.overlay[3])
elif args.auto_resize:
image = utils.resize_input(image)
elif args.auto_resize_crop:
image = utils.resize_crop_input(image)
elif args.auto_rescale:
image = utils.rescale_input(image)

# See if image has the correct shape after preprocessing
if image.shape != (512, 512, 3):
print("Error : image is not 512 x 512, got shape: {}".format(
image.shape), file=sys.stderr)
exit(1)

# Process
if args.n_runs is None or args.n_runs == 1:
result = process(image, gpu_ids, prefs)

if args.overlay:
result = utils.overlay_original_img(original_image, result, args.overlay[0], args.overlay[1],
args.overlay[2], args.overlay[3])

cv2.imwrite(args.output, result)
else:
base_output_filename = utils.strip_file_extension(
args.output, ".png")

def process_one_image(i):
result = process(image, gpu_ids, prefs)

if args.overlay:
result = utils.overlay_original_img(original_image, result, args.overlay[0], args.overlay[1],
args.overlay[2], args.overlay[3])
cv2.imwrite(base_output_filename + "%03d.png" % i, result)

if args.cpu:
pool = ThreadPool(args.n_cores)
pool.map(process_one_image, range(args.n_runs))
pool.close()
pool.join()
else:
for i in range(args.n_runs):
process_one_image(i)
if conf.args['gpu_ids']:
conf.log.info("GAN Processing Will Use GPU IDs: {}".format(conf.args['gpu_ids']))
else:
# Read images
gif_imgs = imageio.mimread(args.input)
print("Total {} frames in the gif!".format(len(gif_imgs)))

# Preprocess
if args.auto_resize:
gif_imgs = [utils.resize_input(img) for img in gif_imgs]
elif args.auto_resize_crop:
gif_imgs = [utils.resize_crop_input(img) for img in gif_imgs]
elif args.auto_rescale:
gif_imgs = [utils.rescale_input(img) for img in gif_imgs]
conf.log.info("GAN Processing Will Use CPU")

# Process
if args.n_runs is None or args.n_runs == 1:
process_gif_wrapper(gif_imgs, args.output if args.output != "output.png" else "output.gif", gpu_ids, prefs, args.n_cores)
else:
base_output_filename = utils.strip_file_extension(args.output,
".gif") if args.output != "output.png" else "output"
for i in range(args.n_runs):
process_gif_wrapper(gif_imgs, base_output_filename + "%03d.gif" % i, gpu_ids, prefs, args.n_cores)

end = time.time()
duration = end - start

# Done
print("Done! We have taken", round(duration, 2), "seconds")
# Processing
start = time.time()
select_processing().run()
conf.log.info("Done! We have taken {} seconds".format(round(time.time() - start, 2)))

# Exit
sys.exit()


# Register Command Handlers
parser.set_defaults(func=main)
gpu_info_parser.set_defaults(func=gpu_info.main)

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:])


def process_gif_wrapper(gif_imgs, filename, gpu_ids, prefs, n_cores):
tmp_dir = tempfile.mkdtemp()
process_gif(gif_imgs, gpu_ids, prefs, tmp_dir, n_cores)
print("Creating gif")
imageio.mimsave(
filename,
[
imageio.imread(os.path.join(tmp_dir, "output_{}.jpg".format(i)))
for i in range(len(gif_imgs))
],
def select_phases():
"""
Select the transformation phases to use following args parameters
:return: <ImageTransform[]> list of image transformation
"""
phases = [DressToCorrect(), CorrectToMask(), MaskToMaskref(),
MaskrefToMaskdet(), MaskdetToMaskfin(), MaskfinToMaskdet()]
if conf.args['overlay']:
phases = [ImageToCrop(), ImageToResized()] + phases + [ImageToOverlay()]
elif conf.args['auto_resize']:
phases = [ImageToResized()] + phases
elif conf.args['auto_resize_crop']:
phases = [ImageToResizedCrop()] + phases
elif conf.args['auto_rescale']:
phases = [ImageToRescale()] + phases
conf.log.debug("Phases to execute : {}".format(phases))
return phases


def select_processing():
"""
Select the processing to use following args parameters
: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)
elif conf.args['n_runs'] != 1:
process = multiple_image_processing(phases, conf.args['n_runs'])
else:
process = simple_image_processing(phases)
conf.log.debug("Process to execute : {}".format(process))
return process


def simple_gif_processing(phases):
"""
Define a simple gif process ready to run
:param phases: <ImageTransform[]> list of image transformation
:return: <SimpleGIFTransform> a gif process run ready
"""
return SimpleGIFTransform(conf.args['input'], phases,
conf.args['output'] if conf.args['output'] != "output.png" else "output.gif")


def multiple_gif_processing(phases, n):
"""
Define a multiple gif process ready to run
:param phases: <ImageTransform[]> list of image transformation
:param n: number of times to process
:return: <MultipleTransform> 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
)
shutil.rmtree(tmp_dir)


def start_rook():
token = os.getenv("ROOKOUT_TOKEN")

if token:
rook.start(token=token)
def simple_image_processing(phases):
"""
Define a simple image process ready to run
:param phases: <ImageTransform[]> list of image transformation
:return: <SimpleImageTransform> a image process run ready
"""
return SimpleImageTransform(conf.args['input'], phases, conf.args['output'])


def multiple_image_processing(phases, n):
"""
Define a multiple image process ready to run
:param phases: <ImageTransform[]> list of image transformation
:param n: number of times to process
:return: <MultipleTransform> a multiple image 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)]
)


if __name__ == "__main__":
freeze_support()
#start_rook()
args.func(args)
# start_rook()
argv.run()

+ 0
- 17
opencv_transform/annotation.py View File

@@ -1,17 +0,0 @@

#Object annotation class:
class BodyPart:
def __init__(self, name, xmin, ymin, xmax, ymax, x, y, w, h):
self.name = name
#Bounding Box:
self.xmin = xmin
self.ymin = ymin
self.xmax = xmax
self.ymax = ymax
#Center:
self.x = x
self.y = y
#Dimensione:
self.w = w
self.h = h

+ 0
- 64
opencv_transform/dress_to_correct.py View File

@@ -1,64 +0,0 @@
import cv2
import math
import numpy as np
import os

# create_correct ===============================================================
# return:
# (<Boolean> True/False), depending on the transformation process
def create_correct(cv_dress):

#Production dir:
return correct_color(cv_dress, 5)

# correct_color ==============================================================================
# return:
# <RGB> image corrected
def correct_color(img, percent):

assert img.shape[2] == 3
assert percent > 0 and percent < 100

half_percent = percent / 200.0

channels = cv2.split(img)

out_channels = []
for channel in channels:
assert len(channel.shape) == 2
# find the low and high precentile values (based on the input percentile)
height, width = channel.shape
vec_size = width * height
flat = channel.reshape(vec_size)

assert len(flat.shape) == 1

flat = np.sort(flat)

n_cols = flat.shape[0]

low_val = flat[math.floor(n_cols * half_percent)]
high_val = flat[math.ceil( n_cols * (1.0 - half_percent))]

# saturate below the low percentile and above the high percentile
thresholded = apply_threshold(channel, low_val, high_val)
# scale the channel
normalized = cv2.normalize(thresholded, thresholded.copy(), 0, 255, cv2.NORM_MINMAX)
out_channels.append(normalized)

return cv2.merge(out_channels)

#Color correction utils
def apply_threshold(matrix, low_value, high_value):
low_mask = matrix < low_value
matrix = apply_mask(matrix, low_mask, low_value)

high_mask = matrix > high_value
matrix = apply_mask(matrix, high_mask, high_value)

return matrix

#Color correction utils
def apply_mask(matrix, mask, fill_value):
masked = np.ma.array(matrix, mask=mask, fill_value=fill_value)
return masked.filled()

+ 0
- 41
opencv_transform/mask_to_maskref.py View File

@@ -1,41 +0,0 @@
import numpy as np
import cv2
import os

###
#
# maskdet_to_maskfin
#
#
###

# create_maskref ===============================================================
# return:
# maskref image
def create_maskref(cv_mask, cv_correct):

#Create a total green image
green = np.zeros((512,512,3), np.uint8)
green[:,:,:] = (0,255,0) # (B, G, R)

#Define the green color filter
f1 = np.asarray([0, 250, 0]) # green color filter
f2 = np.asarray([10, 255, 10])
#From mask, extrapolate only the green mask
green_mask = cv2.inRange(cv_mask, f1, f2) #green is 0

# (OPTIONAL) Apply dilate and open to mask
kernel = np.ones((5,5),np.uint8) #Try change it?
green_mask = cv2.dilate(green_mask, kernel, iterations = 1)
#green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_OPEN, kernel)

# Create an inverted mask
green_mask_inv = cv2.bitwise_not(green_mask)

# Cut correct and green image, using the green_mask & green_mask_inv
res1 = cv2.bitwise_and(cv_correct, cv_correct, mask = green_mask_inv)
res2 = cv2.bitwise_and(green, green, mask = green_mask)

# Compone:
return cv2.add(res1, res2)

+ 0
- 539
opencv_transform/maskdet_to_maskfin.py View File

@@ -1,539 +0,0 @@
import numpy as np
import cv2
import os
import random

#My library:
from opencv_transform.annotation import BodyPart

###
#
# maskdet_to_maskfin
#
# steps:
# 1. Extract annotation
# 1.a: Filter by color
# 1.b: Find ellipses
# 1.c: Filter out ellipses by max size, and max total numbers
# 1.d: Detect Problems
# 1.e: Resolve the problems, or discard the transformation
# 2. With the body list, draw maskfin, using maskref
#
###

# create_maskfin ==============================================================================
# return:
# (<Boolean> True/False), depending on the transformation process
def create_maskfin(maskref, maskdet, prefs):
aur_size = prefs["aursize"]
nip_size = prefs["nipsize"]
tit_size = prefs["titsize"]
vag_size = prefs["vagsize"]
hair_size = prefs["hairsize"]

enable_pubes = (hair_size > 0)

#Create a total green image, in which draw details ellipses
details = np.zeros((512,512,3), np.uint8)
details[:,:,:] = (0,255,0) # (B, G, R)

#Extract body part features:
bodypart_list = extractAnnotations(maskdet, enable_pubes);

#Check if the list is not empty:
if bodypart_list:

#Draw body part in details image:
for obj in bodypart_list:

if obj.w < obj.h:
aMax = int(obj.h/2) #asse maggiore
aMin = int(obj.w/2) #asse minore
angle = 0 #angle
else:
aMax = int(obj.w/2)
aMin = int(obj.h/2)
angle = 90

x = int(obj.x)
y = int(obj.y)

to_int = lambda a, b: int(round(a*float(b)))

aurmax = to_int(aur_size, aMax)
aurmin = to_int(aur_size, aMin)
nipmax = to_int(nip_size, aMax)
nipmin = to_int(nip_size, aMin)
titmax = to_int(tit_size, aMax)
titmin = to_int(tit_size, aMin)
vagmax = to_int(vag_size, aMax)
vagmin = to_int(vag_size, aMin)
hairmax = to_int(hair_size, aMax)
hairmin = to_int(hair_size, aMin)

#Draw ellipse
if obj.name == "tit":
cv2.ellipse(details,(x,y),(titmax,titmin),angle,0,360,(0,205,0),-1) #(0,0,0,50)
elif obj.name == "aur":
cv2.ellipse(details,(x,y),(aurmax,aurmin),angle,0,360,(0,0,255),-1) #red
elif obj.name == "nip":
cv2.ellipse(details,(x,y),(nipmax,nipmin),angle,0,360,(255,255,255),-1) #white
elif obj.name == "belly":
cv2.ellipse(details,(x,y),(aMax,aMin),angle,0,360,(255,0,255),-1) #purple
elif obj.name == "vag":
cv2.ellipse(details,(x,y),(vagmax,vagmin),angle,0,360,(255,0,0),-1) #blue
elif obj.name == "hair":
xmin = x - hairmax
ymin = y - hairmin
xmax = x + hairmax
ymax = y + hairmax
cv2.rectangle(details,(xmin,ymin),(xmax,ymax),(100,100,100),-1)

#Define the green color filter
f1 = np.asarray([0, 250, 0]) # green color filter
f2 = np.asarray([10, 255, 10])

#From maskref, extrapolate only the green mask
green_mask = cv2.bitwise_not(cv2.inRange(maskref, f1, f2)) #green is 0

# Create an inverted mask
green_mask_inv = cv2.bitwise_not(green_mask)

# Cut maskref and detail image, using the green_mask & green_mask_inv
res1 = cv2.bitwise_and(maskref, maskref, mask = green_mask)
res2 = cv2.bitwise_and(details, details, mask = green_mask_inv)

# Compone:
maskfin = cv2.add(res1, res2)
return maskfin

# extractAnnotations ==============================================================================
# input parameter:
# (<string> maskdet_img): relative path of the single maskdet image (es: testimg1/maskdet/1.png)
# return:
# (<BodyPart []> bodypart_list) - for failure/error, return an empty list []
def extractAnnotations(maskdet, enable_pubes):

#Load the image
#image = cv2.imread(maskdet_img)

#Find body part
tits_list = findBodyPart(maskdet, "tit")
aur_list = findBodyPart(maskdet, "aur")
vag_list = findBodyPart(maskdet, "vag")
belly_list = findBodyPart(maskdet, "belly")

#Filter out parts basing on dimension (area and aspect ratio):
aur_list = filterDimParts(aur_list, 100, 1000, 0.5, 3);
tits_list = filterDimParts(tits_list, 1000, 60000, 0.2, 3);
vag_list = filterDimParts(vag_list, 10, 1000, 0.2, 3);
belly_list = filterDimParts(belly_list, 10, 1000, 0.2, 3);

#Filter couple (if parts are > 2, choose only 2)
aur_list = filterCouple(aur_list);
tits_list = filterCouple(tits_list);

#Detect a missing problem:
missing_problem = detectTitAurMissingProblem(tits_list, aur_list) #return a Number (code of the problem)

#Check if problem is SOLVEABLE:
if (missing_problem in [3,6,7,8]):
resolveTitAurMissingProblems(tits_list, aur_list, missing_problem)

#Infer the nips:
nip_list = inferNip(aur_list)

#Infer the hair:
hair_list = inferHair(vag_list, enable_pubes)

#Return a combined list:
return tits_list + aur_list + nip_list + vag_list + hair_list + belly_list

# findBodyPart ==============================================================================
# input parameters:
# (<RGB>image, <string>part_name)
# return
# (<BodyPart[]>list)
def findBodyPart(image, part_name):

bodypart_list = [] #empty BodyPart list

#Get the correct color filter:
if part_name == "tit":
#Use combined color filter
f1 = np.asarray([0, 0, 0]) # tit color filter
f2 = np.asarray([10, 10, 10])
f3 = np.asarray([0, 0, 250]) # aur color filter
f4 = np.asarray([0, 0, 255])
color_mask1 = cv2.inRange(image, f1, f2)
color_mask2 = cv2.inRange(image, f3, f4)
color_mask = cv2.bitwise_or(color_mask1, color_mask2) #combine

elif part_name == "aur":
f1 = np.asarray([0, 0, 250]) # aur color filter
f2 = np.asarray([0, 0, 255])
color_mask = cv2.inRange(image, f1, f2)

elif part_name == "vag":
f1 = np.asarray([250, 0, 0]) # vag filter
f2 = np.asarray([255, 0, 0])
color_mask = cv2.inRange(image, f1, f2)

elif part_name == "belly":
f1 = np.asarray([250, 0, 250]) # belly filter
f2 = np.asarray([255, 0, 255])
color_mask = cv2.inRange(image, f1, f2)

#find contours:
contours, hierarchy = cv2.findContours(color_mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)

#for every contour:
for cnt in contours:

if len(cnt)>5: #at least 5 points to fit ellipse

#(x, y), (MA, ma), angle = cv2.fitEllipse(cnt)
ellipse = cv2.fitEllipse(cnt)

#Fit Result:
x = ellipse[0][0] #center x
y = ellipse[0][1] #center y
angle = ellipse[2] #angle
aMin = ellipse[1][0]; #asse minore
aMax = ellipse[1][1]; #asse maggiore

#Detect direction:
if angle == 0:
h = aMax
w = aMin
else:
h = aMin
w = aMax

#Normalize the belly size:
if part_name == "belly":
if w<15:
w *= 2
if h<15:
h *= 2

#Normalize the vag size:
if part_name == "vag":
if w<15:
w *= 2
if h<15:
h *= 2

#Calculate Bounding Box:
xmin = int(x - (w/2))
xmax = int(x + (w/2))
ymin = int(y - (h/2))
ymax = int(y + (h/2))

bodypart_list.append(BodyPart(part_name, xmin, ymin, xmax, ymax, x, y, w, h ))

return bodypart_list

# filterDimParts ==============================================================================
# input parameters:
# (<BodyPart[]>list, <num> minimum area of part, <num> max area, <num> min aspect ratio, <num> max aspect ratio)
def filterDimParts(bp_list, min_area, max_area, min_ar, max_ar):

b_filt = []

for obj in bp_list:

a = obj.w*obj.h #Object AREA

if ((a > min_area)and(a < max_area)):

ar = obj.w/obj.h #Object ASPECT RATIO

if ((ar>min_ar)and(ar<max_ar)):

b_filt.append(obj)

return b_filt

# filterCouple ==============================================================================
# input parameters:
# (<BodyPart[]>list)
def filterCouple(bp_list):

#Remove exceed parts
if (len(bp_list)>2):

#trovare coppia (a,b) che minimizza bp_list[a].y-bp_list[b].y
min_a = 0
min_b = 1
min_diff = abs(bp_list[min_a].y-bp_list[min_b].y)

for a in range(0,len(bp_list)):
for b in range(0,len(bp_list)):
#TODO: avoid repetition (1,0) (0,1)
if a != b:
diff = abs(bp_list[a].y-bp_list[b].y)
if diff<min_diff:
min_diff = diff
min_a = a
min_b = b
b_filt = []

b_filt.append(bp_list[min_a])
b_filt.append(bp_list[min_b])

return b_filt
else:
#No change
return bp_list



# detectTitAurMissingProblem ==============================================================================
# input parameters:
# (<BodyPart[]> tits list, <BodyPart[]> aur list)
# return
# (<num> problem code)
# TIT | AUR | code | SOLVE? |
# 0 | 0 | 1 | NO |
# 0 | 1 | 2 | NO |
# 0 | 2 | 3 | YES |
# 1 | 0 | 4 | NO |
# 1 | 1 | 5 | NO |
# 1 | 2 | 6 | YES |
# 2 | 0 | 7 | YES |
# 2 | 1 | 8 | YES |
def detectTitAurMissingProblem(tits_list, aur_list):

t_len = len(tits_list)
a_len = len(aur_list)

if (t_len == 0):
if (a_len == 0):
return 1
elif (a_len == 1):
return 2
elif (a_len == 2):
return 3
else:
return -1
elif (t_len == 1):
if (a_len == 0):
return 4
elif (a_len == 1):
return 5
elif (a_len == 2):
return 6
else:
return -1
elif (t_len == 2):
if (a_len == 0):
return 7
elif (a_len == 1):
return 8
else:
return -1
else:
return -1

# resolveTitAurMissingProblems ==============================================================================
# input parameters:
# (<BodyPart[]> tits list, <BodyPart[]> aur list, problem code)
# return
# none
def resolveTitAurMissingProblems(tits_list, aur_list, problem_code):

if problem_code == 3:

random_tit_factor = random.randint(2, 5) #TOTEST

#Add the first tit:
new_w = aur_list[0].w * random_tit_factor #TOTEST
new_x = aur_list[0].x
new_y = aur_list[0].y

xmin = int(new_x - (new_w/2))
xmax = int(new_x + (new_w/2))
ymin = int(new_y - (new_w/2))
ymax = int(new_y + (new_w/2))

tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w ))

#Add the second tit:
new_w = aur_list[1].w * random_tit_factor #TOTEST
new_x = aur_list[1].x
new_y = aur_list[1].y

xmin = int(new_x - (new_w/2))
xmax = int(new_x + (new_w/2))
ymin = int(new_y - (new_w/2))
ymax = int(new_y + (new_w/2))

tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w ))

elif problem_code == 6:

#Find wich aur is full:
d1 = abs(tits_list[0].x - aur_list[0].x)
d2 = abs(tits_list[0].x - aur_list[1].x)

if d1 > d2:
#aur[0] is empty
new_x = aur_list[0].x
new_y = aur_list[0].y
else:
#aur[1] is empty
new_x = aur_list[1].x
new_y = aur_list[1].y

#Calculate Bounding Box:
xmin = int(new_x - (tits_list[0].w/2))
xmax = int(new_x + (tits_list[0].w/2))
ymin = int(new_y - (tits_list[0].w/2))
ymax = int(new_y + (tits_list[0].w/2))

tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, tits_list[0].w, tits_list[0].w ))

elif problem_code == 7:

#Add the first aur:
new_w = tits_list[0].w * random.uniform(0.03, 0.1) #TOTEST
new_x = tits_list[0].x
new_y = tits_list[0].y

xmin = int(new_x - (new_w/2))
xmax = int(new_x + (new_w/2))
ymin = int(new_y - (new_w/2))
ymax = int(new_y + (new_w/2))

aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w ))

#Add the second aur:
new_w = tits_list[1].w * random.uniform(0.03, 0.1) #TOTEST
new_x = tits_list[1].x
new_y = tits_list[1].y

xmin = int(new_x - (new_w/2))
xmax = int(new_x + (new_w/2))
ymin = int(new_y - (new_w/2))
ymax = int(new_y + (new_w/2))

aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w ))

elif problem_code == 8:

#Find wich tit is full:
d1 = abs(aur_list[0].x - tits_list[0].x)
d2 = abs(aur_list[0].x - tits_list[1].x)

if d1 > d2:
#tit[0] is empty
new_x = tits_list[0].x
new_y = tits_list[0].y
else:
#tit[1] is empty
new_x = tits_list[1].x
new_y = tits_list[1].y

#Calculate Bounding Box:
xmin = int(new_x - (aur_list[0].w/2))
xmax = int(new_x + (aur_list[0].w/2))
ymin = int(new_y - (aur_list[0].w/2))
ymax = int(new_y + (aur_list[0].w/2))
aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, aur_list[0].w, aur_list[0].w ))

# detectTitAurPositionProblem ==============================================================================
# input parameters:
# (<BodyPart[]> tits list, <BodyPart[]> aur list)
# return
# (<Boolean> True/False)
def detectTitAurPositionProblem(tits_list, aur_list):

diffTitsX = abs(tits_list[0].x - tits_list[1].x)
if diffTitsX < 40:
print("diffTitsX")
#Tits too narrow (orizontally)
return True

diffTitsY = abs(tits_list[0].y - tits_list[1].y)
if diffTitsY > 120:
#Tits too distanced (vertically)
print("diffTitsY")
return True

diffTitsW = abs(tits_list[0].w - tits_list[1].w)
if ((diffTitsW < 0.1)or(diffTitsW>60)):
print("diffTitsW")
#Tits too equals, or too different (width)
return True

#Check if body position is too low (face not covered by watermark)
if aur_list[0].y > 350: #tits too low
#Calculate the ratio between y and aurs distance
rapp = aur_list[0].y/(abs(aur_list[0].x - aur_list[1].x))
if rapp > 2.8:
print("aurDown")
return True

return False

# inferNip ==============================================================================
# input parameters:
# (<BodyPart[]> aur list)
# return
# (<BodyPart[]> nip list)
def inferNip(aur_list):
nip_list = []

for aur in aur_list:

#Nip rules:
# - circle (w == h)
# - min dim: 5
# - bigger if aur is bigger
nip_dim = int(5 + aur.w*random.uniform(0.03, 0.09))

#center:
x = aur.x
y = aur.y

#Calculate Bounding Box:
xmin = int(x - (nip_dim/2))
xmax = int(x + (nip_dim/2))
ymin = int(y - (nip_dim/2))
ymax = int(y + (nip_dim/2))

nip_list.append(BodyPart("nip", xmin, ymin, xmax, ymax, x, y, nip_dim, nip_dim ))

return nip_list

# inferHair (TOTEST) ==============================================================================
# input parameters:
# (<BodyPart[]> vag list)
# return
# (<BodyPart[]> hair list)
def inferHair(vag_list, enable):
hair_list = []

#70% of chanche to add hair
if enable:

for vag in vag_list:

#Hair rules:
hair_w = vag.w*random.uniform(0.4, 1.5)
hair_h = vag.h*random.uniform(0.4, 1.5)

#center:
x = vag.x
y = vag.y - (hair_h/2) - (vag.h/2)

#Calculate Bounding Box:
xmin = int(x - (hair_w/2))
xmax = int(x + (hair_w/2))
ymin = int(y - (hair_h/2))
ymax = int(y + (hair_h/2))

hair_list.append(BodyPart("hair", xmin, ymin, xmax, ymax, x, y, hair_w, hair_h ))

return hair_list

+ 0
- 29
opencv_transform/nude_to_watermark.py View File

@@ -1,29 +0,0 @@
import cv2
import numpy as np
import os

# create_watermark ===============================================================
# return:
# (<Boolean> True/False), depending on the transformation process
def create_watermark(nude):

# Add alpha channel if missing
if nude.shape[2] < 4:
nude = np.dstack([nude, np.ones((512, 512), dtype="uint8") * 255])

return nude

watermark = cv2.imread("fake.png", cv2.IMREAD_UNCHANGED)

f1 = np.asarray([0, 0, 0, 250]) # red color filter
f2 = np.asarray([255, 255, 255, 255])
mask = cv2.bitwise_not(cv2.inRange(watermark, f1, f2))
mask_inv = cv2.bitwise_not(mask)

res1 = cv2.bitwise_and(nude, nude, mask=mask)
res2 = cv2.bitwise_and(watermark, watermark, mask=mask_inv)
res = cv2.add(res1, res2)

alpha = 0.6
return cv2.addWeighted(res, alpha, nude, 1 - alpha, 0)


+ 41
- 0
processing/__init__.py View File

@@ -0,0 +1,41 @@
import time

from config import Config as conf
from utils import camel_case_to_str


class Process:
"""
Abstract Process Class
"""

def __init__(self):
self.__start = time.time()

def run(self):
self.info_start_run()
self.setup()
self.execute()
self.clean()
self.info_end_run()

def info_start_run(self):
self.__start = time.time()
conf.log.info("Executing {}".format(camel_case_to_str(self.__class__.__name__)))

def info_end_run(self):
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):
pass

def execute(self):
pass

def clean(self):
pass

def __str__(self):
return str(self.__class__.__name__)

+ 58
- 0
processing/gif.py View File

@@ -0,0 +1,58 @@
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: <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
"""
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: <RGB[]> List of final transformed image
"""
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)

+ 91
- 0
processing/image.py View File

@@ -0,0 +1,91 @@
from multiprocessing.pool import ThreadPool

from config import Config as conf
from processing import Process
from utils import read_image, write_image


class SimpleImageTransform(Process):
"""
Simple Image Processing Class
"""

def __init__(self, input_path, phases, output_path):
"""
ProcessImage Constructor
:param input_path: <string> image path to process
:param output_path: <string> image path to write the result.
:param phases: <ImageTransform[]> list of transformation each image
"""
super().__init__()
self.__phases = phases
self.__input_path = input_path
self.__output_path = output_path
self.__image_steps = []

def info_start_run(self):
super().info_start_run()
conf.log.debug("Processing on {}".format(self.__input_path))

def setup(self):
self.__image_steps.append(read_image(self.__input_path))

def execute(self):
"""
Execute all phases on the image
:return: None
"""

for p in self.__phases:
self.__image_steps.append(p.run(*[self.__image_steps[i] for i in p.input_index]))

write_image(self.__image_steps[-1], self.__output_path)
conf.log.debug("{} Image Created ".format(self.__output_path))

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: <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
"""
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: <RGB[]> List of final transformed image
"""

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()

+ 2
- 1
requirements.txt View File

@@ -7,4 +7,5 @@ torch==1.1.0
imageio==2.5.0
python-dotenv==0.10.3
rook==0.1.73
sentry-sdk==0.10.2
sentry-sdk==0.10.2
coloredlogs==10.0

+ 0
- 209
run.py View File

@@ -1,209 +0,0 @@
import cv2
import os

# Import Neural Network Model
from gan import DataLoader, DeepModel, tensor2im

# OpenCv Transform:
from multiprocessing.pool import ThreadPool
from opencv_transform.mask_to_maskref import create_maskref
from opencv_transform.maskdet_to_maskfin import create_maskfin
from opencv_transform.dress_to_correct import create_correct
from opencv_transform.nude_to_watermark import create_watermark

"""
run.py

This script manage the entire transormation.

Transformation happens in 6 phases:
0: dress -> correct [opencv] dress_to_correct
1: correct -> mask: [GAN] correct_to_mask
2: mask -> maskref [opencv] mask_to_maskref
3: maskref -> maskdet [GAN] maskref_to_maskdet
4: maskdet -> maskfin [opencv] maskdet_to_maskfin
5: maskfin -> nude [GAN] maskfin_to_nude
6: nude -> watermark [opencv] nude_to_watermark

"""

phases = [
"dress_to_correct",
"correct_to_mask",
"mask_to_maskref",
"maskref_to_maskdet",
"maskdet_to_maskfin",
"maskfin_to_nude",
"nude_to_watermark",
]


class Options:
# Init options with default values
def __init__(self):

# experiment specifics
self.norm = "batch" # instance normalization or batch normalization
self.use_dropout = False # use dropout for the generator
self.data_type = 32 # Supported data type i.e. 8, 16, 32 bit

# input/output sizes
self.batchSize = 1 # input batch size
self.input_nc = 3 # of input image channels
self.output_nc = 3 # of output image channels

# for setting inputs
# if true, takes images in order to make batches, otherwise takes them randomly
self.serial_batches = True
self.nThreads = (
0
) # threads for loading data. Keep this value at 0! see: https://github.com/pytorch/pytorch/issues/12831
# Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.
self.max_dataset_size = 1

# for generator
self.netG = "global" # selects model to use for netG
self.ngf = 64 # of gen filters in first conv layer
self.n_downsample_global = 4 # number of downsampling layers in netG
self.n_blocks_global = (
9
) # number of residual blocks in the global generator network
self.n_blocks_local = (
0
) # number of residual blocks in the local enhancer network
self.n_local_enhancers = 0 # number of local enhancers to use
# number of epochs that we only train the outmost local enhancer
self.niter_fix_global = 0

# Phase specific options
self.checkpoints_dir = ""
self.dataroot = ""

# Changes options accordlying to actual phase
def updateOptions(self, phase):
directory = os.path.dirname(os.path.realpath(__file__))

if phase == "correct_to_mask":
self.checkpoints_dir = os.path.join(directory, "checkpoints", "cm.lib")

elif phase == "maskref_to_maskdet":
self.checkpoints_dir = os.path.join(directory, "checkpoints", "mm.lib")

elif phase == "maskfin_to_nude":
self.checkpoints_dir = os.path.join(directory, "checkpoints", "mn.lib")


# process(cv_img, mode)
# return:
# watermark image


import sys


def process(cv_img, gpu_ids, prefs):

# InMemory cv2 images:
dress = cv_img
correct = None
mask = None
maskref = None
maskfin = None
maskdet = None
nude = None
watermark = None

print("GPU IDs: " + str(gpu_ids), flush=True)

for index, phase in enumerate(phases):
print("Executing phase: " + phase, flush=True)

# GAN phases:
if (
(phase == "correct_to_mask")
or (phase == "maskref_to_maskdet")
or (phase == "maskfin_to_nude")
):

# Load global option
opt = Options()

# Load custom phase options:
opt.updateOptions(phase)

# Load Data
if phase == "correct_to_mask":
data_loader = DataLoader(opt, correct)
elif phase == "maskref_to_maskdet":
data_loader = DataLoader(opt, maskref)
elif phase == "maskfin_to_nude":
data_loader = DataLoader(opt, maskfin)

dataset = data_loader.load_data()

# Create Model
model = DeepModel()
model.initialize(opt, gpu_ids)

# Run for every image:
for i, data in enumerate(dataset):

generated = model.inference(data["label"], data["inst"])

im = tensor2im(generated.data[0])

# Save Data
if phase == "correct_to_mask":
mask = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
cv2.imwrite("mask.png", mask)

elif phase == "maskref_to_maskdet":
maskdet = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
cv2.imwrite("maskdet.png", maskdet)

elif phase == "maskfin_to_nude":
nude = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
cv2.imwrite("nude.png", nude)

# Correcting:
elif phase == "dress_to_correct":
correct = create_correct(dress)
cv2.imwrite("correct.png", correct)

# mask_ref phase (opencv)
elif phase == "mask_to_maskref":
maskref = create_maskref(mask, correct)
cv2.imwrite("maskref.png", maskref)

# mask_fin phase (opencv)
elif phase == "maskdet_to_maskfin":
maskfin = create_maskfin(maskref, maskdet, prefs)
cv2.imwrite("maskfin.png", maskfin)

# nude_to_watermark phase (opencv)
elif phase == "nude_to_watermark":
watermark = create_watermark(nude)

return watermark


# process(cv_img, mode)
# return:
# gif

def process_gif(gif_imgs, gpu_ids, prefs, tmp_dir, n_cores):
def process_one_image(a):
print("Processing image : {}/{}".format(a[1] + 1, len(gif_imgs)))
img = cv2.resize(a[0], (512, 512))
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imwrite(os.path.join(tmp_dir, "output_{}.jpg".format(a[1])), process(img, gpu_ids, prefs))

print("GPU IDs: " + str(gpu_ids), flush=True)
if gpu_ids is None: # Only multithreading with CPU because threads cause crashes with GPU
pool = ThreadPool(n_cores)
pool.map(process_one_image, zip(gif_imgs, range(len(gif_imgs))))
pool.close()
pool.join()
else:
for x in zip(gif_imgs, range(len(gif_imgs))):
process_one_image(x)

opencv_transform/__init__.py → test.py View File


+ 47
- 0
transform/__init__.py View File

@@ -0,0 +1,47 @@
"""
Image Transformation
"""
import time

from config import Config as conf
from utils import camel_case_to_str


class ImageTransform:
"""
Abstract Image Transformation Class
"""

def __init__(self, input_index=(-1,)):
"""
Image Transformation Class Constructor
:param input_index: index where to take the input (default is -1 for previous transformation)
"""
self.__start = time.time()
self.input_index = input_index

def run(self, *args):
self.__start = time.time()
self.info_start_run()
self.setup(*args)
r = self.execute(*args)
self.clean(*args)
self.info_end_run()
return r

def info_start_run(self):
self.__start = time.time()
conf.log.info("Executing {}".format(camel_case_to_str(self.__class__.__name__)))

def info_end_run(self):
conf.log.debug("{} Done in {} seconds".format(
camel_case_to_str(self.__class__.__name__), round(time.time() - self.__start, 2)))

def setup(self, *args):
pass

def execute(self, *args):
pass

def clean(self, *args):
pass

gan.py → transform/gan/__init__.py View File

@@ -1,13 +1,66 @@
from PIL import Image
import numpy as np
import cv2
import torchvision.transforms as transforms
import torch
import io
import os
import functools
import os
from collections import OrderedDict

import cv2
import numpy as np
import torch
from PIL import Image
from torchvision import transforms as transforms

from config import Config as conf
from transform import ImageTransform


class ImageTransformGAN(ImageTransform):
"""
Abstract GAN Image Transformation Class
"""

def __init__(self, checkpoint, phase):
"""
Abstract GAN Image Transformation Class Constructor
:param checkpoint: <string> path to the checkpoint
:param phase: <string> phase name
"""
super().__init__()
self.__checkpoint = checkpoint
self.__phase = phase
self.__gpu_ids = conf.args["gpu_ids"]

def setup(self, image):
"""
Load Dataset and Model fot the image
:param image: <RGB> image to be transform
:return: None
"""
if self.__gpu_ids:
conf.log.debug("GAN Processing Using GPU IDs: {}".format(self.__gpu_ids))
else:
conf.log.debug("GAN Processing Using CPU")

c = conf()

# Load custom phase options:
data_loader = DataLoader(c, image)
self.__dataset = data_loader.load_data()

# Create Model
self.__model = DeepModel()
self.__model.initialize(c, self.__gpu_ids, self.__checkpoint)

def execute(self, image):
"""
Excute the GAN Transformation the image
:param image: <RGB> image to transform
:return: <RGB> image transformed
"""
for i, data in enumerate(self.__dataset):
generated = self.__model.inference(data["label"], data["inst"])
im = tensor2im(generated.data[0])
mask = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
return mask


class DataLoader:
def __init__(self, opt, cv_img):
@@ -42,7 +95,6 @@ class Dataset(torch.utils.data.Dataset):
self.dataset_size = 1

def __getitem__(self, index):

transform_A = get_transform(self.opt)
A_tensor = transform_A(self.A.convert("RGB"))

@@ -63,9 +115,10 @@ class Dataset(torch.utils.data.Dataset):


class DeepModel(torch.nn.Module):
def initialize(self, opt, gpu_ids):
def initialize(self, opt, gpu_ids, checkpoints_dir):

self.opt = opt
self.checkpoints_dir = checkpoints_dir

if gpu_ids is None:
self.gpu_ids = []
@@ -104,7 +157,7 @@ class DeepModel(torch.nn.Module):
# helper loading function that can be used by subclasses
def __load_network(self, network):

save_path = os.path.join(self.opt.checkpoints_dir)
save_path = os.path.join(self.checkpoints_dir)

state_dict = torch.load(save_path)

@@ -120,7 +173,7 @@ class DeepModel(torch.nn.Module):
network.load_state_dict(new_state_dict)

def __encode_input(
self, label_map, inst_map=None, real_image=None, feat_map=None, infer=False
self, label_map, inst_map=None, real_image=None, feat_map=None, infer=False
):
if len(self.gpu_ids) > 0:
input_label = label_map.data.cuda() # GPU
@@ -138,17 +191,17 @@ class DeepModel(torch.nn.Module):
m.bias.data.fill_(0)

def __define_G(
self,
input_nc,
output_nc,
ngf,
netG,
n_downsample_global=3,
n_blocks_global=9,
n_local_enhancers=1,
n_blocks_local=3,
norm="instance",
gpu_ids=[],
self,
input_nc,
output_nc,
ngf,
netG,
n_downsample_global=3,
n_blocks_global=9,
n_local_enhancers=1,
n_blocks_local=3,
norm="instance",
gpu_ids=[],
):
norm_layer = self.__get_norm_layer(norm_type=norm)

@@ -187,14 +240,14 @@ class DeepModel(torch.nn.Module):

class GlobalGenerator(torch.nn.Module):
def __init__(
self,
input_nc,
output_nc,
ngf=64,
n_downsampling=3,
n_blocks=9,
norm_layer=torch.nn.BatchNorm2d,
padding_type="reflect",
self,
input_nc,
output_nc,
ngf=64,
n_downsampling=3,
n_blocks=9,
norm_layer=torch.nn.BatchNorm2d,
padding_type="reflect",
):
assert n_blocks >= 0
super(GlobalGenerator, self).__init__()
@@ -264,12 +317,12 @@ class GlobalGenerator(torch.nn.Module):

class ResnetBlock(torch.nn.Module):
def __init__(
self,
dim,
padding_type,
norm_layer,
activation=torch.nn.ReLU(True),
use_dropout=False,
self,
dim,
padding_type,
norm_layer,
activation=torch.nn.ReLU(True),
use_dropout=False,
):
super(ResnetBlock, self).__init__()
self.conv_block = self.__build_conv_block(
@@ -277,7 +330,7 @@ class ResnetBlock(torch.nn.Module):
)

def __build_conv_block(
self, dim, padding_type, norm_layer, activation, use_dropout
self, dim, padding_type, norm_layer, activation, use_dropout
):
conv_block = []
p = 0

+ 38
- 0
transform/gan/mask.py View File

@@ -0,0 +1,38 @@
from transform.gan import ImageTransformGAN
from config import Config as opt


class CorrectToMask(ImageTransformGAN):