Browse Source

Added --color-transfer

tags/v1.2.5
PommeDroid 10 months ago
parent
commit
4f791c63d9
6 changed files with 221 additions and 9 deletions
  1. 6
    0
      argv.py
  2. 7
    6
      checkpoints.py
  3. 5
    1
      main.py
  4. 184
    0
      third/opencv/color_transfer.py
  5. 18
    0
      transform/opencv/correct.py
  6. 1
    2
      utils.py

+ 6
- 0
argv.py View File

@@ -275,6 +275,12 @@ class ArgvParser:

return type_func

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

ArgvParser.parser.add_argument(
"-s",
"--steps",

+ 7
- 6
checkpoints.py View File

@@ -5,7 +5,7 @@ import sys
import tempfile

from config import Config as conf
from utils import setup_log, dll_file, unzip
from utils import setup_log, dl_file, unzip


def main(_):
@@ -25,16 +25,17 @@ def download(_):

try:
conf.log.info("Downloading {}".format(cdn_url))
dll_file(conf.checkpoints_cdn.format(conf.checkpoints_version), temp_zip)
dl_file(conf.checkpoints_cdn.format(conf.checkpoints_version), temp_zip)

conf.log.info("Extracting {}".format(temp_zip))
unzip(temp_zip, conf.args['checkpoints'])

conf.log.info("Moving Checkpoints To Final location")

[(lambda a: os.remove(a) and shutil.move(a, os.path.abspath(conf.args['checkpoints'])))(x)
for x in (os.path.join(conf.args['checkpoints'], 'checkpoints', y) for y in ("cm.lib", "mm.lib", "mn.lib"))]
conf.log.info("Moving Checkpoints To Final Location")

for c in ("cm.lib", "mm.lib", "mn.lib"):
if os.path.isfile(os.path.join(conf.args['checkpoints'], c)):
os.remove(os.path.join(conf.args['checkpoints'], c))
shutil.move(os.path.join(conf.args['checkpoints'], 'checkpoints', c), conf.args['checkpoints'])
shutil.rmtree(os.path.join(conf.args['checkpoints'], 'checkpoints'))

except Exception as e:

+ 5
- 1
main.py View File

@@ -11,7 +11,7 @@ from utils import setup_log, check_shape
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
from transform.opencv.correct import DressToCorrect, ColorTransfer
from transform.opencv.mask import MaskToMaskref, MaskdetToMaskfin


@@ -83,6 +83,10 @@ def select_phases():
check_shape(conf.args['input'])
else:
conf.log.warn('Image Size Requirements Unchecked.')

if conf.args['color_transfer']:
phases = add_head(phases, ColorTransfer)

return phases



+ 184
- 0
third/opencv/color_transfer.py View File

@@ -0,0 +1,184 @@
"""
The MIT License (MIT)

Copyright (c) 2014 Adrian Rosebrock, http://www.pyimagesearch.com

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""
# import the necessary packages
import numpy as np
import cv2


def color_transfer(source, target, clip=True, preserve_paper=True):
"""
Transfers the color distribution from the source to the target
image using the mean and standard deviations of the L*a*b*
color space.

This implementation is (loosely) based on to the "Color Transfer
between Images" paper by Reinhard et al., 2001.

Parameters:
-------
source: NumPy array
OpenCV image in BGR color space (the source image)
target: NumPy array
OpenCV image in BGR color space (the target image)
clip: Should components of L*a*b* image be scaled by np.clip before
converting back to BGR color space?
If False then components will be min-max scaled appropriately.
Clipping will keep target image brightness truer to the input.
Scaling will adjust image brightness to avoid washed out portions
in the resulting color transfer that can be caused by clipping.
preserve_paper: Should color transfer strictly follow methodology
layed out in original paper? The method does not always produce
aesthetically pleasing results.
If False then L*a*b* components will scaled using the reciprocal of
the scaling factor proposed in the paper. This method seems to produce
more consistently aesthetically pleasing results

Returns:
-------
transfer: NumPy array
OpenCV image (w, h, 3) NumPy array (uint8)
"""
# convert the images from the RGB to L*ab* color space, being
# sure to utilizing the floating point data type (note: OpenCV
# expects floats to be 32-bit, so use that instead of 64-bit)
source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype("float32")
target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype("float32")

# compute color statistics for the source and target images
(lMeanSrc, lStdSrc, aMeanSrc, aStdSrc, bMeanSrc, bStdSrc) = image_stats(source)
(lMeanTar, lStdTar, aMeanTar, aStdTar, bMeanTar, bStdTar) = image_stats(target)

# subtract the means from the target image
(l, a, b) = cv2.split(target)
l -= lMeanTar
a -= aMeanTar
b -= bMeanTar

if preserve_paper:
# scale by the standard deviations using paper proposed factor
l = (lStdTar / lStdSrc) * l
a = (aStdTar / aStdSrc) * a
b = (bStdTar / bStdSrc) * b
else:
# scale by the standard deviations using reciprocal of paper proposed factor
l = (lStdSrc / lStdTar) * l
a = (aStdSrc / aStdTar) * a
b = (bStdSrc / bStdTar) * b

# add in the source mean
l += lMeanSrc
a += aMeanSrc
b += bMeanSrc

# clip/scale the pixel intensities to [0, 255] if they fall
# outside this range
l = _scale_array(l, clip=clip)
a = _scale_array(a, clip=clip)
b = _scale_array(b, clip=clip)

# merge the channels together and convert back to the RGB color
# space, being sure to utilize the 8-bit unsigned integer data
# type
transfer = cv2.merge([l, a, b])
transfer = cv2.cvtColor(transfer.astype("uint8"), cv2.COLOR_LAB2BGR)

# return the color transferred image
return transfer


def image_stats(image):
"""
Parameters:
-------
image: NumPy array
OpenCV image in L*a*b* color space

Returns:
-------
Tuple of mean and standard deviations for the L*, a*, and b*
channels, respectively
"""
# compute the mean and standard deviation of each channel
(l, a, b) = cv2.split(image)
(lMean, lStd) = (l.mean(), l.std())
(aMean, aStd) = (a.mean(), a.std())
(bMean, bStd) = (b.mean(), b.std())

# return the color statistics
return (lMean, lStd, aMean, aStd, bMean, bStd)


def _min_max_scale(arr, new_range=(0, 255)):
"""
Perform min-max scaling to a NumPy array

Parameters:
-------
arr: NumPy array to be scaled to [new_min, new_max] range
new_range: tuple of form (min, max) specifying range of
transformed array

Returns:
-------
NumPy array that has been scaled to be in
[new_range[0], new_range[1]] range
"""
# get array's current min and max
mn = arr.min()
mx = arr.max()

# check if scaling needs to be done to be in new_range
if mn < new_range[0] or mx > new_range[1]:
# perform min-max scaling
scaled = (new_range[1] - new_range[0]) * (arr - mn) / (mx - mn) + new_range[0]
else:
# return array if already in range
scaled = arr

return scaled


def _scale_array(arr, clip=True):
"""
Trim NumPy array values to be in [0, 255] range with option of
clipping or scaling.

Parameters:
-------
arr: array to be trimmed to [0, 255] range
clip: should array be scaled by np.clip? if False then input
array will be min-max scaled to range
[max([arr.min(), 0]), min([arr.max(), 255])]

Returns:
-------
NumPy array that has been scaled to be in [0, 255] range
"""
if clip:
scaled = np.clip(arr, 0, 255)
else:
scale_range = (max([arr.min(), 0]), min([arr.max(), 255]))
scaled = _min_max_scale(arr, new_range=scale_range)

return scaled

+ 18
- 0
transform/opencv/correct.py View File

@@ -2,6 +2,7 @@ import cv2
import math
import numpy as np

from third.opencv.color_transfer import color_transfer
from transform.opencv import ImageTransformOpenCV


@@ -86,3 +87,20 @@ class DressToCorrect(ImageTransformOpenCV):
"""
masked = np.ma.array(matrix, mask=mask, fill_value=fill_value)
return masked.filled()


class ColorTransfer(ImageTransformOpenCV):
"""
ColorTransfer [OPENCV]
"""
def __init__(self, input_index=(0, -1), args=None):
super().__init__(input_index=input_index, args=args)

def execute(self, img, img_target):
"""
Transfers the color distribution from the source to the target
:param img: <RGB> Image source
:param img_target: <RGB> Image target
:return: <RGB> Color transfer image
"""
return color_transfer(img, img_target, clip=True, preserve_paper=False)

+ 1
- 2
utils.py View File

@@ -132,7 +132,7 @@ def json_to_argv(data):
return l


def dll_file(url, file_path):
def dl_file(url, file_path):
"""
Download a file
:param url: <string> url of the file to download
@@ -160,7 +160,6 @@ def dll_file(url, file_path):
done = int(50 * dl / total_length)
print("[{}{}]".format('=' * done, ' ' * (50 - done)), end="\r")
print(" "*80, end="\r")
conf.log.info("{} Downloaded".format(url,))
return file_path



Loading…
Cancel
Save