diff --git a/core/imagelib/color_transfer.py b/core/imagelib/color_transfer.py index c605920..aaadd88 100644 --- a/core/imagelib/color_transfer.py +++ b/core/imagelib/color_transfer.py @@ -1,7 +1,7 @@ import cv2 import numpy as np from numpy import linalg as npla -from random import random, shuffle, choice +import random from scipy.stats import special_ortho_group import scipy as sp @@ -92,7 +92,7 @@ def color_transfer_mkl(x0, x1): def color_transfer_idt(i0, i1, bins=256, n_rot=20): import scipy.stats - + relaxation = 1 / n_rot h,w,c = i0.shape h1,w1,c1 = i1.shape @@ -371,23 +371,24 @@ def color_transfer(ct_mode, img_src, img_trg): # imported from faceswap -def color_augmentation(img): +def color_augmentation(img, seed=None): """ Color adjust RGB image """ face = img face = np.clip(face*255.0, 0, 255).astype(np.uint8) - face = random_clahe(face) - face = random_lab(face) + face = random_clahe(face, seed) + face = random_lab(face, seed) img[:, :, :3] = face return (face / 255.0).astype(np.float32) -def random_lab(image): +def random_lab(image, seed=None): """ Perform random color/lightness adjustment in L*a*b* colorspace """ + random.seed(seed) amount_l = 30 / 100 amount_ab = 8 / 100 - randoms = [(random() * amount_l * 2) - amount_l, # L adjust - (random() * amount_ab * 2) - amount_ab, # A adjust - (random() * amount_ab * 2) - amount_ab] # B adjust + randoms = [(random.random() * amount_l * 2) - amount_l, # L adjust + (random.random() * amount_ab * 2) - amount_ab, # A adjust + (random.random() * amount_ab * 2) - amount_ab] # B adjust image = cv2.cvtColor( # pylint:disable=no-member image, cv2.COLOR_BGR2LAB).astype("float32") / 255.0 # pylint:disable=no-member @@ -400,15 +401,16 @@ def random_lab(image): cv2.COLOR_LAB2BGR) # pylint:disable=no-member return image -def random_clahe(image): +def random_clahe(image, seed=None): """ Randomly perform Contrast Limited Adaptive Histogram Equalization """ - contrast_random = random() + random.seed(seed) + contrast_random = random.random() if contrast_random > 50 / 100: return image # base_contrast = image.shape[0] // 128 base_contrast = 1 # testing because it breaks on small sizes - grid_base = random() * 4 + grid_base = random.random() * 4 contrast_adjustment = int(grid_base * (base_contrast / 2)) grid_size = base_contrast + contrast_adjustment @@ -416,4 +418,4 @@ def random_clahe(image): tileGridSize=(grid_size, grid_size)) for chan in range(3): image[:, :, chan] = clahe.apply(image[:, :, chan]) - return image \ No newline at end of file + return image diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index 2e514dc..a2f1e0e 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -204,7 +204,7 @@ class SampleProcessor(object): # Apply random color transfer if ct_mode is not None and ct_sample is not None or ct_mode == 'fs-aug': if ct_mode == 'fs-aug': - img = imagelib.color_augmentation(img) + img = imagelib.color_augmentation(img, sample_rnd_seed) else: if ct_sample_bgr is None: ct_sample_bgr = ct_sample.load_bgr()