mirror of
https://github.com/iperov/DeepFaceLive
synced 2025-07-16 10:03:42 -07:00
refactoring
This commit is contained in:
parent
6b6b6b2d16
commit
64116844f2
9 changed files with 360 additions and 74 deletions
Binary file not shown.
|
@ -135,7 +135,7 @@ class FaceAlignerWorker(BackendWorker):
|
||||||
exclude_moving_parts=state.exclude_moving_parts,
|
exclude_moving_parts=state.exclude_moving_parts,
|
||||||
head_yaw=head_yaw,
|
head_yaw=head_yaw,
|
||||||
x_offset=state.x_offset,
|
x_offset=state.x_offset,
|
||||||
y_offset=state.y_offset)
|
y_offset=state.y_offset-0.08)
|
||||||
|
|
||||||
fsi.face_align_image_name = f'{frame_image_name}_{face_id}_aligned'
|
fsi.face_align_image_name = f'{frame_image_name}_{face_id}_aligned'
|
||||||
fsi.image_to_align_uni_mat = uni_mat
|
fsi.image_to_align_uni_mat = uni_mat
|
||||||
|
|
|
@ -164,7 +164,7 @@ class FaceMarkerWorker(BackendWorker):
|
||||||
if fsi.face_urect is not None:
|
if fsi.face_urect is not None:
|
||||||
# Cut the face to feed to the face marker
|
# Cut the face to feed to the face marker
|
||||||
face_image, face_uni_mat = fsi.face_urect.cut(frame_image, marker_state.marker_coverage, 256 if is_opencv_lbf else \
|
face_image, face_uni_mat = fsi.face_urect.cut(frame_image, marker_state.marker_coverage, 256 if is_opencv_lbf else \
|
||||||
192 if is_google_facemesh else 0 )
|
192 if is_google_facemesh else 0 )
|
||||||
_,H,W,_ = ImageProcessor(face_image).get_dims()
|
_,H,W,_ = ImageProcessor(face_image).get_dims()
|
||||||
|
|
||||||
if is_opencv_lbf:
|
if is_opencv_lbf:
|
||||||
|
|
|
@ -233,7 +233,7 @@ class FaceMergerWorker(BackendWorker):
|
||||||
|
|
||||||
# Combine face mask
|
# Combine face mask
|
||||||
face_mask = ImageProcessor(face_mask).erode_blur(state.face_mask_erode, state.face_mask_blur, fade_to_border=True).get_image('HWC')
|
face_mask = ImageProcessor(face_mask).erode_blur(state.face_mask_erode, state.face_mask_blur, fade_to_border=True).get_image('HWC')
|
||||||
frame_face_mask = ImageProcessor(face_mask).warpAffine(aligned_to_source_uni_mat, frame_width, frame_height).clip2( (1.0/255.0), 0.0, 1.0, 1.0).get_image('HWC')
|
frame_face_mask = ImageProcessor(face_mask).warp_affine(aligned_to_source_uni_mat, frame_width, frame_height).clip2( (1.0/255.0), 0.0, 1.0, 1.0).get_image('HWC')
|
||||||
|
|
||||||
face_swap_img = ImageProcessor(face_swap_img).to_ufloat32().get_image('HWC')
|
face_swap_img = ImageProcessor(face_swap_img).to_ufloat32().get_image('HWC')
|
||||||
|
|
||||||
|
@ -241,7 +241,7 @@ class FaceMergerWorker(BackendWorker):
|
||||||
face_align_img = ImageProcessor(face_align_img).to_ufloat32().get_image('HWC')
|
face_align_img = ImageProcessor(face_align_img).to_ufloat32().get_image('HWC')
|
||||||
face_swap_img = lib_ct.rct(face_swap_img, face_align_img, target_mask=face_mask, source_mask=face_mask)
|
face_swap_img = lib_ct.rct(face_swap_img, face_align_img, target_mask=face_mask, source_mask=face_mask)
|
||||||
|
|
||||||
frame_face_swap_img = ImageProcessor(face_swap_img).warpAffine(aligned_to_source_uni_mat, frame_width, frame_height, interpolation=interpolation).get_image('HWC')
|
frame_face_swap_img = ImageProcessor(face_swap_img).warp_affine(aligned_to_source_uni_mat, frame_width, frame_height, interpolation=interpolation).get_image('HWC')
|
||||||
|
|
||||||
# Combine final frame
|
# Combine final frame
|
||||||
opacity = np.float32(state.face_opacity)
|
opacity = np.float32(state.face_opacity)
|
||||||
|
|
|
@ -251,10 +251,10 @@ class FaceSwapperWorker(BackendWorker):
|
||||||
|
|
||||||
fai_ip = ImageProcessor(face_align_image)
|
fai_ip = ImageProcessor(face_align_image)
|
||||||
if model_state.presharpen_amount != 0:
|
if model_state.presharpen_amount != 0:
|
||||||
fai_ip.sharpen(factor=model_state.presharpen_amount)
|
fai_ip.gaussian_sharpen(sigma=1.0, power=model_state.presharpen_amount)
|
||||||
|
|
||||||
if pre_gamma_red != 1.0 or pre_gamma_green != 1.0 or pre_gamma_blue != 1.0:
|
if pre_gamma_red != 1.0 or pre_gamma_green != 1.0 or pre_gamma_blue != 1.0:
|
||||||
fai_ip.adjust_gamma(pre_gamma_red, pre_gamma_green, pre_gamma_blue)
|
fai_ip.gamma(pre_gamma_red, pre_gamma_green, pre_gamma_blue)
|
||||||
face_align_image = fai_ip.get_image('HWC')
|
face_align_image = fai_ip.get_image('HWC')
|
||||||
|
|
||||||
celeb_face, celeb_face_mask_img, face_align_mask_img = dfm_model.convert(face_align_image, morph_factor=model_state.morph_factor)
|
celeb_face, celeb_face_mask_img, face_align_mask_img = dfm_model.convert(face_align_image, morph_factor=model_state.morph_factor)
|
||||||
|
|
|
@ -77,8 +77,8 @@ class FrameAdjusterWorker(BackendWorker):
|
||||||
|
|
||||||
if frame_image is not None:
|
if frame_image is not None:
|
||||||
frame_image_ip = ImageProcessor(frame_image)
|
frame_image_ip = ImageProcessor(frame_image)
|
||||||
frame_image_ip.median_blur(5, state.median_blur_per / 100.0 )
|
frame_image_ip.median_blur(5, opacity=state.median_blur_per / 100.0 )
|
||||||
frame_image_ip.degrade_resize( state.degrade_bicubic_per / 100.0, interpolation=ImageProcessor.Interpolation.CUBIC)
|
frame_image_ip.reresize( state.degrade_bicubic_per / 100.0, interpolation=ImageProcessor.Interpolation.CUBIC)
|
||||||
|
|
||||||
frame_image = frame_image_ip.get_image('HWC')
|
frame_image = frame_image_ip.get_image('HWC')
|
||||||
bcd.set_image(frame_image_name, frame_image)
|
bcd.set_image(frame_image_name, frame_image)
|
||||||
|
|
20
main.py
20
main.py
|
@ -4,6 +4,7 @@ import platform
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from xlib import appargs as lib_appargs
|
from xlib import appargs as lib_appargs
|
||||||
|
from xlib import os as lib_os
|
||||||
|
|
||||||
# onnxruntime==1.8.0 requires CUDA_PATH_V11_2, but 1.8.1 don't
|
# onnxruntime==1.8.0 requires CUDA_PATH_V11_2, but 1.8.1 don't
|
||||||
# keep the code if they return that behaviour
|
# keep the code if they return that behaviour
|
||||||
|
@ -54,7 +55,7 @@ def main():
|
||||||
|
|
||||||
def run_extract_FaceSynthetics(args):
|
def run_extract_FaceSynthetics(args):
|
||||||
from scripts import dev
|
from scripts import dev
|
||||||
|
|
||||||
inputdir_path = Path(args.input_dir)
|
inputdir_path = Path(args.input_dir)
|
||||||
faceset_path = Path(args.faceset_path)
|
faceset_path = Path(args.faceset_path)
|
||||||
|
|
||||||
|
@ -69,23 +70,14 @@ def main():
|
||||||
train_parsers = train_parser.add_subparsers()
|
train_parsers = train_parser.add_subparsers()
|
||||||
|
|
||||||
def train_FaceAligner(args):
|
def train_FaceAligner(args):
|
||||||
faceset_path = Path(args.faceset_path)
|
lib_os.set_process_priority(lib_os.ProcessPriority.IDLE)
|
||||||
|
from apps.trainers.FaceAligner.FaceAlignerTrainerApp import FaceAlignerTrainerApp
|
||||||
from apps.trainers.FaceAligner.FaceAlignerTrainer import FaceAlignerTrainer
|
FaceAlignerTrainerApp(workspace_path=Path(args.workspace_dir), faceset_path=Path(args.faceset_path))
|
||||||
FaceAlignerTrainer(faceset_path=faceset_path).run()
|
|
||||||
|
|
||||||
p = train_parsers.add_parser('FaceAligner')
|
p = train_parsers.add_parser('FaceAligner')
|
||||||
|
p.add_argument('--workspace-dir', default=None, action=fixPathAction, help="Workspace directory.")
|
||||||
p.add_argument('--faceset-path', default=None, action=fixPathAction, help=".dfs path")
|
p.add_argument('--faceset-path', default=None, action=fixPathAction, help=".dfs path")
|
||||||
p.set_defaults(func=train_FaceAligner)
|
p.set_defaults(func=train_FaceAligner)
|
||||||
|
|
||||||
def train_CTSOT(args):
|
|
||||||
from apps.trainers.CTSOT.CTSOTTrainerApp import CTSOTTrainerApp
|
|
||||||
CTSOTTrainerApp(workspace_path=Path(args.workspace_dir), faceset_path=Path(args.faceset_path))
|
|
||||||
|
|
||||||
p = train_parsers.add_parser('CTSOT')
|
|
||||||
p.add_argument('--workspace-dir', default=None, action=fixPathAction, help="Workspace directory.")
|
|
||||||
p.add_argument('--faceset-path', default=None, action=fixPathAction, help=".dfs faceset path")
|
|
||||||
p.set_defaults(func=train_CTSOT)
|
|
||||||
|
|
||||||
def bad_args(arguments):
|
def bad_args(arguments):
|
||||||
parser.print_help()
|
parser.print_help()
|
||||||
|
|
|
@ -102,6 +102,8 @@ class FLandmarks2D(IState):
|
||||||
r = max(xrt[0], xrb[0])
|
r = max(xrt[0], xrb[0])
|
||||||
b = max(xlb[1], xrb[1])
|
b = max(xlb[1], xrb[1])
|
||||||
return FRect.from_ltrb( (l,t,r,b) )
|
return FRect.from_ltrb( (l,t,r,b) )
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
def calc_cut(self, h_w, coverage : float, output_size : int,
|
def calc_cut(self, h_w, coverage : float, output_size : int,
|
||||||
exclude_moving_parts : bool = False,
|
exclude_moving_parts : bool = False,
|
||||||
|
@ -143,19 +145,18 @@ class FLandmarks2D(IState):
|
||||||
bt_diag_vec = (g_p[1]-g_p[3]).astype(np.float32)
|
bt_diag_vec = (g_p[1]-g_p[3]).astype(np.float32)
|
||||||
bt_diag_vec /= npla.norm(bt_diag_vec)
|
bt_diag_vec /= npla.norm(bt_diag_vec)
|
||||||
|
|
||||||
# calc modifier of diagonal vectors for scale and coverage value
|
# calc modifier of diagonal vectors for coverage value
|
||||||
scale = 1.0
|
mod = npla.norm(g_p[0]-g_p[2])*(coverage*0.5)
|
||||||
mod = (1.0 / scale)* ( npla.norm(g_p[0]-g_p[2])*( coverage * 0.5) )
|
|
||||||
|
|
||||||
# adjust vertical offset to cover more forehead
|
|
||||||
h_vec = (g_p[1]-g_p[0]).astype(np.float32)
|
|
||||||
v_vec = (g_p[3]-g_p[0]).astype(np.float32)
|
|
||||||
|
|
||||||
if head_yaw is not None:
|
if head_yaw is not None:
|
||||||
# Damp near zero
|
# Damp near zero
|
||||||
x_offset += -(head_yaw * np.abs(np.tanh(head_yaw*2)) ) * 0.5
|
x_offset += -(head_yaw * np.abs(np.tanh(head_yaw*2)) ) * 0.5
|
||||||
|
|
||||||
g_c += h_vec*x_offset + v_vec*(y_offset-0.08)
|
# adjust vertical offset to cover more forehead
|
||||||
|
h_vec = (g_p[1]-g_p[0]).astype(np.float32)
|
||||||
|
v_vec = (g_p[3]-g_p[0]).astype(np.float32)
|
||||||
|
|
||||||
|
g_c += h_vec*x_offset + v_vec*y_offset
|
||||||
|
|
||||||
l_t = np.array( [ g_c - tb_diag_vec*mod,
|
l_t = np.array( [ g_c - tb_diag_vec*mod,
|
||||||
g_c + bt_diag_vec*mod,
|
g_c + bt_diag_vec*mod,
|
||||||
|
@ -174,7 +175,7 @@ class FLandmarks2D(IState):
|
||||||
exclude_moving_parts : bool = False,
|
exclude_moving_parts : bool = False,
|
||||||
head_yaw : float = None,
|
head_yaw : float = None,
|
||||||
x_offset : float = 0,
|
x_offset : float = 0,
|
||||||
y_offset : float = 0) -> Tuple[Affine2DMat, Affine2DUniMat]:
|
y_offset : float = 0) -> Tuple[np.ndarray, Affine2DUniMat]:
|
||||||
"""
|
"""
|
||||||
Cut the face to square of output_size from img using landmarks with given parameters
|
Cut the face to square of output_size from img using landmarks with given parameters
|
||||||
|
|
||||||
|
|
|
@ -39,7 +39,7 @@ class ImageProcessor:
|
||||||
"""
|
"""
|
||||||
"""
|
"""
|
||||||
ip = ImageProcessor.__new__(ImageProcessor)
|
ip = ImageProcessor.__new__(ImageProcessor)
|
||||||
ip._img = self._img
|
ip._img = self._img.copy()
|
||||||
return ip
|
return ip
|
||||||
|
|
||||||
def get_dims(self) -> Tuple[int,int,int,int]:
|
def get_dims(self) -> Tuple[int,int,int,int]:
|
||||||
|
@ -53,18 +53,24 @@ class ImageProcessor:
|
||||||
def get_dtype(self):
|
def get_dtype(self):
|
||||||
return self._img.dtype
|
return self._img.dtype
|
||||||
|
|
||||||
def adjust_gamma(self, red : float, green : float, blue : float) -> 'ImageProcessor':
|
def gamma(self, red : float, green : float, blue : float, mask=None) -> 'ImageProcessor':
|
||||||
dtype = self.get_dtype()
|
dtype = self.get_dtype()
|
||||||
self.to_ufloat32()
|
self.to_ufloat32()
|
||||||
img = self._img
|
img = orig_img = self._img
|
||||||
np.power(img, np.array([1.0 / blue, 1.0 / green, 1.0 / red], np.float32), out=img)
|
|
||||||
|
img = np.power(img, np.array([1.0 / blue, 1.0 / green, 1.0 / red], np.float32) )
|
||||||
np.clip(img, 0, 1.0, out=img)
|
np.clip(img, 0, 1.0, out=img)
|
||||||
|
|
||||||
|
if mask is not None:
|
||||||
|
mask = self._check_normalize_mask(mask)
|
||||||
|
img = ne.evaluate('orig_img*(1-mask) + img*mask')
|
||||||
|
|
||||||
self._img = img
|
self._img = img
|
||||||
self.to_dtype(dtype)
|
self.to_dtype(dtype)
|
||||||
return self
|
return self
|
||||||
|
|
||||||
|
|
||||||
def apply(self, func) -> 'ImageProcessor':
|
def apply(self, func, mask=None) -> 'ImageProcessor':
|
||||||
"""
|
"""
|
||||||
apply your own function on internal image
|
apply your own function on internal image
|
||||||
|
|
||||||
|
@ -76,12 +82,16 @@ class ImageProcessor:
|
||||||
|
|
||||||
.apply( lambda img: img-[102,127,63] )
|
.apply( lambda img: img-[102,127,63] )
|
||||||
"""
|
"""
|
||||||
img = self._img
|
img = orig_img = self._img
|
||||||
dtype = img.dtype
|
img = func(img).astype(orig_img.dtype)
|
||||||
new_img = func(self._img).astype(dtype)
|
if img.ndim != 4:
|
||||||
if new_img.ndim != 4:
|
|
||||||
raise Exception('func used in ImageProcessor.apply changed format of image')
|
raise Exception('func used in ImageProcessor.apply changed format of image')
|
||||||
self._img = new_img
|
|
||||||
|
if mask is not None:
|
||||||
|
mask = self._check_normalize_mask(mask)
|
||||||
|
img = ne.evaluate('orig_img*(1-mask) + img*mask').astype(orig_img.dtype)
|
||||||
|
|
||||||
|
self._img = img
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def fit_in (self, TW = None, TH = None, pad_to_target : bool = False, allow_upscale : bool = False, interpolation : 'ImageProcessor.Interpolation' = None) -> float:
|
def fit_in (self, TW = None, TH = None, pad_to_target : bool = False, allow_upscale : bool = False, interpolation : 'ImageProcessor.Interpolation' = None) -> float:
|
||||||
|
@ -147,7 +157,7 @@ class ImageProcessor:
|
||||||
img[h] = high_val
|
img[h] = high_val
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def degrade_resize(self, power : float, interpolation : 'ImageProcessor.Interpolation' = None) -> 'ImageProcessor':
|
def reresize(self, power : float, interpolation : 'ImageProcessor.Interpolation' = None, mask = None) -> 'ImageProcessor':
|
||||||
"""
|
"""
|
||||||
|
|
||||||
power float 0 .. 1.0
|
power float 0 .. 1.0
|
||||||
|
@ -159,7 +169,7 @@ class ImageProcessor:
|
||||||
if interpolation is None:
|
if interpolation is None:
|
||||||
interpolation = ImageProcessor.Interpolation.LINEAR
|
interpolation = ImageProcessor.Interpolation.LINEAR
|
||||||
|
|
||||||
img = self._img
|
img = orig_img = self._img
|
||||||
|
|
||||||
N,H,W,C = img.shape
|
N,H,W,C = img.shape
|
||||||
W_lr = max(4, int(W*(1.0-power)))
|
W_lr = max(4, int(W*(1.0-power)))
|
||||||
|
@ -168,41 +178,196 @@ class ImageProcessor:
|
||||||
img = cv2.resize (img, (W_lr,H_lr), interpolation=_cv_inter[interpolation])
|
img = cv2.resize (img, (W_lr,H_lr), interpolation=_cv_inter[interpolation])
|
||||||
img = cv2.resize (img, (W,H) , interpolation=_cv_inter[interpolation])
|
img = cv2.resize (img, (W,H) , interpolation=_cv_inter[interpolation])
|
||||||
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
|
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
|
||||||
|
|
||||||
|
if mask is not None:
|
||||||
|
mask = self._check_normalize_mask(mask)
|
||||||
|
img = ne.evaluate('orig_img*(1-mask) + img*mask').astype(orig_img.dtype)
|
||||||
|
|
||||||
self._img = img
|
self._img = img
|
||||||
return self
|
return self
|
||||||
|
|
||||||
|
def box_sharpen(self, size : int, power : float, mask = None) -> 'ImageProcessor':
|
||||||
def median_blur(self, size : int, power : float) -> 'ImageProcessor':
|
|
||||||
"""
|
"""
|
||||||
size int median kernel size
|
size int kernel size
|
||||||
|
|
||||||
power float 0 .. 1.0
|
power float 0 .. 1.0 (or higher)
|
||||||
"""
|
"""
|
||||||
power = min(1, max(0, power))
|
power = max(0, power)
|
||||||
if power == 0:
|
if power == 0:
|
||||||
return self
|
return self
|
||||||
|
|
||||||
|
if size % 2 == 0:
|
||||||
|
size += 1
|
||||||
|
|
||||||
|
dtype = self.get_dtype()
|
||||||
|
self.to_ufloat32()
|
||||||
|
|
||||||
|
img = orig_img = self._img
|
||||||
|
N,H,W,C = img.shape
|
||||||
|
|
||||||
|
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
|
||||||
|
|
||||||
|
kernel = np.zeros( (size, size), dtype=np.float32)
|
||||||
|
kernel[ size//2, size//2] = 1.0
|
||||||
|
box_filter = np.ones( (size, size), dtype=np.float32) / (size**2)
|
||||||
|
kernel = kernel + (kernel - box_filter) * (power)
|
||||||
|
img = cv2.filter2D(img, -1, kernel)
|
||||||
|
img = np.clip(img, 0, 1, out=img)
|
||||||
|
|
||||||
|
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
|
||||||
|
|
||||||
|
if mask is not None:
|
||||||
|
mask = self._check_normalize_mask(mask)
|
||||||
|
img = ne.evaluate('orig_img*(1-mask) + img*mask')
|
||||||
|
|
||||||
|
self._img = img
|
||||||
|
self.to_dtype(dtype)
|
||||||
|
return self
|
||||||
|
|
||||||
|
def gaussian_sharpen(self, sigma : float, power : float, mask = None) -> 'ImageProcessor':
|
||||||
|
"""
|
||||||
|
sigma float
|
||||||
|
|
||||||
|
power float 0 .. 1.0 and higher
|
||||||
|
"""
|
||||||
|
sigma = max(0, sigma)
|
||||||
|
if sigma == 0:
|
||||||
|
return self
|
||||||
|
|
||||||
dtype = self.get_dtype()
|
dtype = self.get_dtype()
|
||||||
self.to_ufloat32()
|
self.to_ufloat32()
|
||||||
|
|
||||||
img = self._img
|
img = orig_img = self._img
|
||||||
N,H,W,C = img.shape
|
N,H,W,C = img.shape
|
||||||
|
|
||||||
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
|
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
|
||||||
|
|
||||||
img_blur = cv2.medianBlur(img, size)
|
img = cv2.addWeighted(img, 1.0 + power,
|
||||||
img = ne.evaluate('img*(1.0-power) + img_blur*power')
|
cv2.GaussianBlur(img, (0, 0), sigma), -power, 0)
|
||||||
|
img = np.clip(img, 0, 1, out=img)
|
||||||
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
|
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
|
||||||
|
|
||||||
|
if mask is not None:
|
||||||
|
mask = self._check_normalize_mask(mask)
|
||||||
|
img = ne.evaluate('orig_img*(1-mask) + img*mask')
|
||||||
|
|
||||||
self._img = img
|
self._img = img
|
||||||
|
|
||||||
self.to_dtype(dtype)
|
self.to_dtype(dtype)
|
||||||
return self
|
return self
|
||||||
|
|
||||||
|
def gaussian_blur(self, sigma : float, opacity : float = 1.0, mask = None) -> 'ImageProcessor':
|
||||||
|
"""
|
||||||
|
sigma float
|
||||||
|
|
||||||
|
opacity float 0 .. 1.0
|
||||||
|
"""
|
||||||
|
sigma = max(0, sigma)
|
||||||
|
if sigma == 0:
|
||||||
|
return self
|
||||||
|
opacity = np.float32( min(1, max(0, opacity)) )
|
||||||
|
if opacity == 0:
|
||||||
|
return self
|
||||||
|
|
||||||
|
dtype = self.get_dtype()
|
||||||
|
self.to_ufloat32()
|
||||||
|
|
||||||
|
img = orig_img = self._img
|
||||||
|
N,H,W,C = img.shape
|
||||||
|
|
||||||
|
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
|
||||||
|
|
||||||
|
img_blur = cv2.GaussianBlur(img, (0,0), sigma)
|
||||||
|
f32_1 = np.float32(1.0)
|
||||||
|
img = ne.evaluate('img*(f32_1-opacity) + img_blur*opacity')
|
||||||
|
|
||||||
|
img = np.clip(img, 0, 1, out=img)
|
||||||
|
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
|
||||||
|
|
||||||
|
if mask is not None:
|
||||||
|
mask = self._check_normalize_mask(mask)
|
||||||
|
img = ne.evaluate('orig_img*(1-mask) + img*mask')
|
||||||
|
|
||||||
|
self._img = img
|
||||||
|
|
||||||
|
self.to_dtype(dtype)
|
||||||
|
return self
|
||||||
|
|
||||||
|
def median_blur(self, size : int, opacity : float, mask = None) -> 'ImageProcessor':
|
||||||
|
"""
|
||||||
|
size int median kernel size
|
||||||
|
|
||||||
|
opacity float 0 .. 1.0
|
||||||
|
"""
|
||||||
|
opacity = min(1, max(0, opacity))
|
||||||
|
if opacity == 0:
|
||||||
|
return self
|
||||||
|
|
||||||
|
dtype = self.get_dtype()
|
||||||
|
self.to_ufloat32()
|
||||||
|
|
||||||
|
img = orig_img = self._img
|
||||||
|
N,H,W,C = img.shape
|
||||||
|
|
||||||
|
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
|
||||||
|
|
||||||
|
img_blur = cv2.medianBlur(img, size)
|
||||||
|
f32_1 = np.float32(1.0)
|
||||||
|
img = ne.evaluate('img*(f32_1-opacity) + img_blur*opacity')
|
||||||
|
img = np.clip(img, 0, 1, out=img)
|
||||||
|
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
|
||||||
|
|
||||||
|
if mask is not None:
|
||||||
|
mask = self._check_normalize_mask(mask)
|
||||||
|
img = ne.evaluate('orig_img*(1-mask) + img*mask')
|
||||||
|
|
||||||
|
self._img = img
|
||||||
|
|
||||||
|
self.to_dtype(dtype)
|
||||||
|
return self
|
||||||
|
|
||||||
|
def motion_blur( self, size, angle, mask=None ):
|
||||||
|
"""
|
||||||
|
size [1..]
|
||||||
|
|
||||||
|
angle degrees
|
||||||
|
|
||||||
|
mask H,W
|
||||||
|
H,W,C
|
||||||
|
N,H,W,C int/float 0-1 will be applied
|
||||||
|
"""
|
||||||
|
if size % 2 == 0:
|
||||||
|
size += 1
|
||||||
|
|
||||||
|
dtype = self.get_dtype()
|
||||||
|
self.to_ufloat32()
|
||||||
|
|
||||||
|
img = orig_img = self._img
|
||||||
|
N,H,W,C = img.shape
|
||||||
|
|
||||||
|
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
|
||||||
|
|
||||||
|
k = np.zeros((size, size), dtype=np.float32)
|
||||||
|
k[ (size-1)// 2 , :] = np.ones(size, dtype=np.float32)
|
||||||
|
k = cv2.warpAffine(k, cv2.getRotationMatrix2D( (size / 2 -0.5 , size / 2 -0.5 ) , angle, 1.0), (size, size) )
|
||||||
|
k = k * ( 1.0 / np.sum(k) )
|
||||||
|
|
||||||
|
img = cv2.filter2D(img, -1, k)
|
||||||
|
img = np.clip(img, 0, 1, out=img)
|
||||||
|
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
|
||||||
|
|
||||||
|
if mask is not None:
|
||||||
|
mask = self._check_normalize_mask(mask)
|
||||||
|
img = ne.evaluate('orig_img*(1-mask) + img*mask')
|
||||||
|
|
||||||
|
self._img = img
|
||||||
|
self.to_dtype(dtype)
|
||||||
|
return self
|
||||||
|
|
||||||
|
|
||||||
def erode_blur(self, erode : int, blur : int, fade_to_border : bool = False) -> 'ImageProcessor':
|
def erode_blur(self, erode : int, blur : int, fade_to_border : bool = False) -> 'ImageProcessor':
|
||||||
"""
|
"""
|
||||||
apply erode and blur to the image
|
apply erode and blur to the mask image
|
||||||
|
|
||||||
erode int != 0
|
erode int != 0
|
||||||
blur int > 0
|
blur int > 0
|
||||||
|
@ -244,7 +409,126 @@ class ImageProcessor:
|
||||||
|
|
||||||
self._img = img
|
self._img = img
|
||||||
return self
|
return self
|
||||||
|
|
||||||
|
def levels(self, in_bwg_out_bw, mask = None) -> 'ImageProcessor':
|
||||||
|
"""
|
||||||
|
in_bwg_out_bw ( [N],[C], 5)
|
||||||
|
optional per channel/batch input black,white,gamma and out black,white floats
|
||||||
|
|
||||||
|
in black = [0.0 .. 1.0] default:0.0
|
||||||
|
in white = [0.0 .. 1.0] default:1.0
|
||||||
|
in gamma = [0.0 .. 2.0++] default:1.0
|
||||||
|
|
||||||
|
out black = [0.0 .. 1.0] default:0.0
|
||||||
|
out white = [0.0 .. 1.0] default:1.0
|
||||||
|
"""
|
||||||
|
dtype = self.get_dtype()
|
||||||
|
self.to_ufloat32()
|
||||||
|
|
||||||
|
img = orig_img = self._img
|
||||||
|
N,H,W,C = img.shape
|
||||||
|
|
||||||
|
v = np.array(in_bwg_out_bw, np.float32)
|
||||||
|
|
||||||
|
if v.ndim == 1:
|
||||||
|
v = v[None,None,...]
|
||||||
|
v = np.tile(v, (N,C,1))
|
||||||
|
elif v.ndim == 2:
|
||||||
|
v = v[None,...]
|
||||||
|
v = np.tile(v, (N,1,1))
|
||||||
|
elif v.ndim > 3:
|
||||||
|
raise ValueError('in_bwg_out_bw.ndim > 3')
|
||||||
|
|
||||||
|
VN, VC, VD = v.shape
|
||||||
|
if N != VN or C != VC or VD != 5:
|
||||||
|
raise ValueError('wrong in_bwg_out_bw size. Must have 5 floats at last dim.')
|
||||||
|
|
||||||
|
v = v[:,None,None,:,:]
|
||||||
|
|
||||||
|
img = np.clip( (img - v[...,0]) / (v[...,1] - v[...,0]), 0, 1 )
|
||||||
|
|
||||||
|
img = ( img ** (1/v[...,2]) ) * (v[...,4] - v[...,3]) + v[...,3]
|
||||||
|
img = np.clip(img, 0, 1, out=img)
|
||||||
|
|
||||||
|
if mask is not None:
|
||||||
|
mask = self._check_normalize_mask(mask)
|
||||||
|
img = ne.evaluate('orig_img*(1-mask) + img*mask')
|
||||||
|
|
||||||
|
self._img = img
|
||||||
|
self.to_dtype(dtype)
|
||||||
|
return self
|
||||||
|
|
||||||
|
def hsv(self, h_diff : float, s_diff : float, v_diff : float, mask = None) -> 'ImageProcessor':
|
||||||
|
"""
|
||||||
|
apply HSV modification for BGR image
|
||||||
|
|
||||||
|
h_diff = [-360.0 .. 360.0]
|
||||||
|
s_diff = [-1.0 .. 1.0]
|
||||||
|
s_diff = [-1.0 .. 1.0]
|
||||||
|
"""
|
||||||
|
dtype = self.get_dtype()
|
||||||
|
self.to_ufloat32()
|
||||||
|
|
||||||
|
img = orig_img = self._img
|
||||||
|
N,H,W,C = img.shape
|
||||||
|
if C != 3:
|
||||||
|
raise Exception('Image channels must be == 3')
|
||||||
|
|
||||||
|
img = img.reshape( (N*H,W,C) )
|
||||||
|
|
||||||
|
h, s, v = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
|
||||||
|
h = ( h + h_diff ) % 360
|
||||||
|
|
||||||
|
s += s_diff
|
||||||
|
np.clip (s, 0, 1, out=s )
|
||||||
|
|
||||||
|
v += v_diff
|
||||||
|
np.clip (v, 0, 1, out=v )
|
||||||
|
|
||||||
|
img = np.clip( cv2.cvtColor(cv2.merge([h, s, v]), cv2.COLOR_HSV2BGR) , 0, 1 )
|
||||||
|
img = img.reshape( (N,H,W,C) )
|
||||||
|
|
||||||
|
if mask is not None:
|
||||||
|
mask = self._check_normalize_mask(mask)
|
||||||
|
img = ne.evaluate('orig_img*(1-mask) + img*mask')
|
||||||
|
|
||||||
|
self._img = img
|
||||||
|
self.to_dtype(dtype)
|
||||||
|
return self
|
||||||
|
|
||||||
|
def jpeg_recompress(self, quality : int, mask = None ) -> 'ImageProcessor':
|
||||||
|
"""
|
||||||
|
quality 0-100
|
||||||
|
"""
|
||||||
|
dtype = self.get_dtype()
|
||||||
|
self.to_uint8()
|
||||||
|
|
||||||
|
img = orig_img = self._img
|
||||||
|
_,_,_,C = img.shape
|
||||||
|
if C != 3:
|
||||||
|
raise Exception('Image channels must be == 3')
|
||||||
|
|
||||||
|
new_imgs = []
|
||||||
|
for x in img:
|
||||||
|
ret, result = cv2.imencode('.jpg', x, [int(cv2.IMWRITE_JPEG_QUALITY), quality] )
|
||||||
|
if not ret:
|
||||||
|
raise Exception('unable to compress jpeg')
|
||||||
|
x = cv2.imdecode(result, flags=cv2.IMREAD_UNCHANGED)
|
||||||
|
|
||||||
|
new_imgs.append(x)
|
||||||
|
|
||||||
|
img = np.array(new_imgs)
|
||||||
|
|
||||||
|
|
||||||
|
if mask is not None:
|
||||||
|
mask = self._check_normalize_mask(mask)
|
||||||
|
img = ne.evaluate('orig_img*(1-mask) + img*mask').astype(np.uint8)
|
||||||
|
|
||||||
|
self._img = img
|
||||||
|
self.to_dtype(dtype)
|
||||||
|
|
||||||
|
return self
|
||||||
|
|
||||||
def rotate90(self) -> 'ImageProcessor':
|
def rotate90(self) -> 'ImageProcessor':
|
||||||
self._img = np.rot90(self._img, k=1, axes=(1,2) )
|
self._img = np.rot90(self._img, k=1, axes=(1,2) )
|
||||||
return self
|
return self
|
||||||
|
@ -305,18 +589,6 @@ class ImageProcessor:
|
||||||
self._img = img
|
self._img = img
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def sharpen(self, factor : float, kernel_size=3) -> 'ImageProcessor':
|
|
||||||
img = self._img
|
|
||||||
|
|
||||||
N,H,W,C = img.shape
|
|
||||||
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
|
|
||||||
blur = cv2.GaussianBlur(img, (kernel_size, kernel_size) , 0)
|
|
||||||
img = cv2.addWeighted(img, 1.0 + (0.5 * factor), blur, -(0.5 * factor), 0)
|
|
||||||
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
|
|
||||||
|
|
||||||
self._img = img
|
|
||||||
return self
|
|
||||||
|
|
||||||
def get_image(self, format) -> np.ndarray:
|
def get_image(self, format) -> np.ndarray:
|
||||||
"""
|
"""
|
||||||
returns image with desired format
|
returns image with desired format
|
||||||
|
@ -395,7 +667,7 @@ class ImageProcessor:
|
||||||
|
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def resize(self, size : Tuple, interpolation : 'ImageProcessor.Interpolation' = None, new_ip=False ) -> 'ImageProcessor':
|
def resize(self, size : Tuple, interpolation : 'ImageProcessor.Interpolation' = None ) -> 'ImageProcessor':
|
||||||
"""
|
"""
|
||||||
resize to (W,H)
|
resize to (W,H)
|
||||||
"""
|
"""
|
||||||
|
@ -411,14 +683,11 @@ class ImageProcessor:
|
||||||
img = cv2.resize (img, (TW, TH), interpolation=_cv_inter[interpolation])
|
img = cv2.resize (img, (TW, TH), interpolation=_cv_inter[interpolation])
|
||||||
img = img.reshape( (TH,TW,N,C) ).transpose( (2,0,1,3) )
|
img = img.reshape( (TH,TW,N,C) ).transpose( (2,0,1,3) )
|
||||||
|
|
||||||
if new_ip:
|
|
||||||
return ImageProcessor(img)
|
|
||||||
|
|
||||||
self._img = img
|
self._img = img
|
||||||
|
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def warpAffine(self, mat, out_width, out_height, interpolation : 'ImageProcessor.Interpolation' = None ) -> 'ImageProcessor':
|
def warp_affine(self, mat, out_width, out_height, interpolation : 'ImageProcessor.Interpolation' = None ) -> 'ImageProcessor':
|
||||||
"""
|
"""
|
||||||
img HWC
|
img HWC
|
||||||
"""
|
"""
|
||||||
|
@ -489,12 +758,36 @@ class ImageProcessor:
|
||||||
self._img = img.astype(np.uint8, copy=False)
|
self._img = img.astype(np.uint8, copy=False)
|
||||||
return self
|
return self
|
||||||
|
|
||||||
|
def _check_normalize_mask(self, mask : np.ndarray):
|
||||||
|
N,H,W,C = self._img.shape
|
||||||
|
|
||||||
|
if mask.ndim == 2:
|
||||||
|
mask = mask[None,...,None]
|
||||||
|
elif mask.ndim == 3:
|
||||||
|
mask = mask[None,...]
|
||||||
|
|
||||||
|
if mask.ndim != 4:
|
||||||
|
raise ValueError('mask must have ndim == 4')
|
||||||
|
|
||||||
|
MN, MH, MW, MC = mask.shape
|
||||||
|
if H != MH or W != MW:
|
||||||
|
raise ValueError('mask H,W, mismatch')
|
||||||
|
|
||||||
|
if MN != 1 and N != MN:
|
||||||
|
raise ValueError(f'mask N dim must be 1 or == {N}')
|
||||||
|
if MC != 1 and C != MC:
|
||||||
|
raise ValueError(f'mask C dim must be 1 or == {C}')
|
||||||
|
|
||||||
|
return mask
|
||||||
|
|
||||||
class Interpolation(IntEnum):
|
class Interpolation(IntEnum):
|
||||||
LINEAR = 0
|
NEAREST = 0,
|
||||||
CUBIC = 1
|
LINEAR = 1
|
||||||
|
CUBIC = 2,
|
||||||
LANCZOS4 = 4
|
LANCZOS4 = 4
|
||||||
|
|
||||||
_cv_inter = { ImageProcessor.Interpolation.LINEAR : cv2.INTER_LINEAR,
|
_cv_inter = { ImageProcessor.Interpolation.NEAREST : cv2.INTER_NEAREST,
|
||||||
|
ImageProcessor.Interpolation.LINEAR : cv2.INTER_LINEAR,
|
||||||
ImageProcessor.Interpolation.CUBIC : cv2.INTER_CUBIC,
|
ImageProcessor.Interpolation.CUBIC : cv2.INTER_CUBIC,
|
||||||
ImageProcessor.Interpolation.LANCZOS4 : cv2.INTER_LANCZOS4,
|
ImageProcessor.Interpolation.LANCZOS4 : cv2.INTER_LANCZOS4,
|
||||||
}
|
}
|
Loading…
Add table
Add a link
Reference in a new issue