ConverterMasked: removed default transfercolor,

added Apply color transfer to predicted face - modes rct / lct

models/ConverterMasked.py

@@ -34,7 +34,6 @@ class ConverterMasked(ConverterBase):
         self.output_size = output_size
         self.face_type = face_type 
         self.clip_hborder_mask_per = clip_hborder_mask_per
-        self.TFLabConverter = None
         
         mode = input_int ("Choose mode: (1) overlay, (2) hist match, (3) hist match bw, (4) seamless, (5) seamless hist match, (6) raw. Default - %d : " % (default_mode) , default_mode)
         
@@ -72,9 +71,9 @@ class ConverterMasked(ConverterBase):
                 self.seamless_erode_mask_modifier = np.clip ( input_int ("Choose seamless erode mask modifier [-100..100] (skip:0) : ", 0), -100, 100)
           
         self.output_face_scale = np.clip ( 1.0 + input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0)*0.01, 0.5, 1.5)
-        
+        self.color_transfer_mode = input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct'])
+
         if self.mode != 'raw':
-            self.transfercolor = input_bool("Transfer color from dst face to converted final face? (y/n skip:n) : ", False)
             self.final_image_color_degrade_power = np.clip (  input_int ("Degrade color power of final image [0..100] (skip:0) : ", 0), 0, 100)
             self.alpha = input_bool("Export png with alpha channel? (y/n skip:n) : ", False)
             
@@ -215,6 +214,29 @@ class ConverterMasked(ConverterBase):
                     if debug:
                         debugs += [img_mask_blurry_aaa.copy()]
                         
+                    if self.color_transfer_mode is not None:
+                        if self.color_transfer_mode == 'rct':
+                            if debug:
+                                debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]
+                          
+                            prd_face_bgr = image_utils.reinhard_color_transfer ( np.clip( (prd_face_bgr*255).astype(np.uint8), 0, 255),
+                                                                                 np.clip( (dst_face_bgr*255).astype(np.uint8), 0, 255) )
+                            prd_face_bgr = np.clip( prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0)
+
+                            if debug:
+                                debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]
+                          
+                            
+                        elif self.color_transfer_mode == 'lct':
+                            if debug:
+                                debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]
+
+                            prd_face_bgr = image_utils.linear_color_transfer (prd_face_bgr, dst_face_bgr)
+                            prd_face_bgr = np.clip( prd_face_bgr, 0.0, 1.0)
+                            
+                            if debug:
+                                debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]
+ 
                     if self.mode == 'hist-match-bw':
                         prd_face_bgr = cv2.cvtColor(prd_face_bgr, cv2.COLOR_BGR2GRAY)
                         prd_face_bgr = np.repeat( np.expand_dims (prd_face_bgr, -1), (3,), -1 )
@@ -238,9 +260,10 @@ class ConverterMasked(ConverterBase):
                             
                     if self.mode == 'hist-match-bw':
                         prd_face_bgr = prd_face_bgr.astype(dtype=np.float32)
-
+                        
                     out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT )
-            
+                    out_img = np.clip(out_img, 0.0, 1.0)
+                    
                     if debug:
                         debugs += [out_img.copy()]
 
@@ -265,16 +288,7 @@ class ConverterMasked(ConverterBase):
                         new_out_face_bgr = image_utils.color_hist_match(out_face_bgr, dst_face_bgr, self.hist_match_threshold)                
                         new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT )
                         out_img =  np.clip( img_bgr*(1-img_mask_blurry_aaa) + (new_out*img_mask_blurry_aaa) , 0, 1.0 )
-                        
+                    
-                    if self.transfercolor:  
-                        if self.TFLabConverter is None:
-                            self.TFLabConverter = image_utils.TFLabConverter() 
-                            
-                        img_lab_l, img_lab_a, img_lab_b = np.split ( self.TFLabConverter.bgr2lab (img_bgr), 3, axis=-1 )
-                        out_img_lab_l, out_img_lab_a, out_img_lab_b = np.split ( self.TFLabConverter.bgr2lab (out_img), 3, axis=-1 )      
-                        
-                        out_img = self.TFLabConverter.lab2bgr ( np.concatenate([out_img_lab_l, img_lab_a, img_lab_b], axis=-1) )
-         
                     if self.final_image_color_degrade_power != 0:
                         if debug:
                             debugs += [out_img.copy()]                    

utils/image_utils.py

@@ -7,6 +7,151 @@ from scipy.spatial import Delaunay
 from PIL import Image, ImageDraw, ImageFont
 from nnlib import nnlib
 
+def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, source_mask=None, target_mask=None):
+	"""
+	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(np.float32)
+	target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype(np.float32)       
+    
+	# compute color statistics for the source and target images
+	src_input = source if source_mask is None else source*source_mask
+	tgt_input = target if target_mask is None else target*target_mask
+	(lMeanSrc, lStdSrc, aMeanSrc, aStdSrc, bMeanSrc, bStdSrc) = lab_image_stats(src_input)
+	(lMeanTar, lStdTar, aMeanTar, aStdTar, bMeanTar, bStdTar) = lab_image_stats(tgt_input)
+
+	# 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(np.uint8), cv2.COLOR_LAB2BGR)
+	
+	# return the color transferred image
+	return transfer
+
+def linear_color_transfer(target_img, source_img, mode='pca', eps=1e-5):
+    '''
+    Matches the colour distribution of the target image to that of the source image
+    using a linear transform.
+    Images are expected to be of form (w,h,c) and float in [0,1].
+    Modes are chol, pca or sym for different choices of basis.
+    '''
+    mu_t = target_img.mean(0).mean(0)
+    t = target_img - mu_t
+    t = t.transpose(2,0,1).reshape(3,-1)
+    Ct = t.dot(t.T) / t.shape[1] + eps * np.eye(t.shape[0])
+    mu_s = source_img.mean(0).mean(0)
+    s = source_img - mu_s
+    s = s.transpose(2,0,1).reshape(3,-1)
+    Cs = s.dot(s.T) / s.shape[1] + eps * np.eye(s.shape[0])
+    if mode == 'chol':
+        chol_t = np.linalg.cholesky(Ct)
+        chol_s = np.linalg.cholesky(Cs)
+        ts = chol_s.dot(np.linalg.inv(chol_t)).dot(t)
+    if mode == 'pca':
+        eva_t, eve_t = np.linalg.eigh(Ct)
+        Qt = eve_t.dot(np.sqrt(np.diag(eva_t))).dot(eve_t.T)
+        eva_s, eve_s = np.linalg.eigh(Cs)
+        Qs = eve_s.dot(np.sqrt(np.diag(eva_s))).dot(eve_s.T)
+        ts = Qs.dot(np.linalg.inv(Qt)).dot(t)
+    if mode == 'sym':
+        eva_t, eve_t = np.linalg.eigh(Ct)
+        Qt = eve_t.dot(np.sqrt(np.diag(eva_t))).dot(eve_t.T)
+        Qt_Cs_Qt = Qt.dot(Cs).dot(Qt)
+        eva_QtCsQt, eve_QtCsQt = np.linalg.eigh(Qt_Cs_Qt)
+        QtCsQt = eve_QtCsQt.dot(np.sqrt(np.diag(eva_QtCsQt))).dot(eve_QtCsQt.T)
+        ts = np.linalg.inv(Qt).dot(QtCsQt).dot(np.linalg.inv(Qt)).dot(t)
+    matched_img = ts.reshape(*target_img.transpose(2,0,1).shape).transpose(1,2,0)
+    matched_img += mu_s
+    matched_img[matched_img>1] = 1
+    matched_img[matched_img<0] = 0
+    return matched_img
+    
+def lab_image_stats(image):
+    # 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 _scale_array(arr, clip=True):
+    if clip:
+        return np.clip(arr, 0, 255)
+
+    mn = arr.min()
+    mx = arr.max()
+    scale_range = (max([mn, 0]), min([mx, 255]))
+   
+    if mn < scale_range[0] or mx > scale_range[1]:
+        return (scale_range[1] - scale_range[0]) * (arr - mn) / (mx - mn) + scale_range[0]
+
+    return arr
+    
 def channel_hist_match(source, template, hist_match_threshold=255, mask=None):
     # Code borrowed from:
     # https://stackoverflow.com/questions/32655686/histogram-matching-of-two-images-in-python-2-x
@@ -154,7 +299,6 @@ def morph_by_points (image, sp, dp):
 
     result_image = np.zeros(image.shape, dtype = image.dtype)
 
-    
     for tri in Delaunay(dp).simplices:                                    
         morphTriangle(result_image, image, sp[tri], dp[tri])