removing CuPy. FaceMerger now works with any OpenCL1.2-compatible GPU.

This commit is contained in:
iperov 2021-09-30 18:31:11 +04:00
parent c2ba2bab9d
commit 4fe2da23c0
6 changed files with 134 additions and 292 deletions

View file

@ -7,33 +7,16 @@ import numpy as np
class ImageProcessor:
"""
Generic image processor for numpy or cupy images
Generic image processor for numpy images
arguments
img np.ndarray|
cp.ndarray
HW (2 ndim)
img np.ndarray HW (2 ndim)
HWC (3 ndim)
NHWC (4 ndim)
for cupy you should set device before using ImageProcessor
"""
def __init__(self, img : Union[np.ndarray,'cp.ndarray'], copy=False):
if img.__class__ == np.ndarray:
self._xp = np
import scipy
import scipy.ndimage
self._sp = scipy
if copy:
img = img.copy()
else:
import cupy as cp # BUG eats 1.8Gb paging file per process, so import on demand
import cupyx.scipy.ndimage
self._xp = cp
self._sp = cupyx.scipy
def __init__(self, img : np.ndarray, copy=False):
ndim = img.ndim
if ndim not in [2,3,4]:
raise ValueError(f'img.ndim must be 2,3,4, not {ndim}.')
@ -55,8 +38,6 @@ class ImageProcessor:
"""
ip = ImageProcessor.__new__(ImageProcessor)
ip._img = self._img
ip._xp = self._xp
ip._sp = self._sp
return ip
def get_dims(self) -> Tuple[int,int,int,int]:
@ -73,16 +54,11 @@ class ImageProcessor:
def adjust_gamma(self, red : float, green : float, blue : float) -> 'ImageProcessor':
dtype = self.get_dtype()
self.to_ufloat32()
xp, img = self._xp , self._img,
xp.power(img, xp.array([1.0 / blue, 1.0 / green, 1.0 / red], xp.float32), out=img)
xp.clip(img, 0, 1.0, out=img)
img = self._img
np.power(img, np.array([1.0 / blue, 1.0 / green, 1.0 / red], np.float32), out=img)
np.clip(img, 0, 1.0, out=img)
self._img = img
self.to_dtype(dtype)
return self
@ -124,7 +100,6 @@ class ImageProcessor:
"""
#if interpolation is None:
# interpolation = ImageProcessor.Interpolation.LINEAR
xp, sp = self._xp, self._sp
img = self._img
N,H,W,C = img.shape
@ -146,12 +121,7 @@ class ImageProcessor:
if scale != 1.0:
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
if self._xp == np:
img = cv2.resize (img, ( int(W*scale), int(H*scale) ), interpolation=ImageProcessor.Interpolation.LINEAR)
else:
img = sp.ndimage.zoom(img, (scale, scale, 1.0), order=1)
img = cv2.resize (img, ( int(W*scale), int(H*scale) ), interpolation=ImageProcessor.Interpolation.LINEAR)
H,W = img.shape[0:2]
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
@ -159,14 +129,13 @@ class ImageProcessor:
w_pad = (TW-W) if TW is not None else 0
h_pad = (TH-H) if TH is not None else 0
if w_pad != 0 or h_pad != 0:
img = xp.pad(img, ( (0,0), (0,h_pad), (0,w_pad), (0,0) ))
img = np.pad(img, ( (0,0), (0,h_pad), (0,w_pad), (0,0) ))
self._img = img
return scale
def clip(self, min, max) -> 'ImageProcessor':
xp = self._xp
xp.clip(self._img, min, max, out=self._img)
np.clip(self._img, min, max, out=self._img)
return self
def clip2(self, low_check, low_val, high_check, high_val) -> 'ImageProcessor':
@ -188,22 +157,14 @@ class ImageProcessor:
if interpolation is None:
interpolation = ImageProcessor.Interpolation.LINEAR
xp, sp, img = self._xp, self._sp, self._img
img = self._img
N,H,W,C = img.shape
W_lr = max(4, int(W*(1.0-power)))
H_lr = max(4, int(H*(1.0-power)))
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
if xp == np:
W_lr = max(4, int(W*(1.0-power)))
H_lr = max(4, int(H*(1.0-power)))
img = cv2.resize (img, (W_lr,H_lr), interpolation=_cv_inter[interpolation])
img = cv2.resize (img, (W,H) , interpolation=_cv_inter[interpolation])
else:
W_lr = max(4, round(W*(1.0-power)))
H_lr = max(4, round(H*(1.0-power)))
img = sp.ndimage.zoom(img, (H_lr/H, W_lr/W, 1), order=_scipy_order[interpolation])
img = sp.ndimage.zoom(img, (H/img.shape[0], W/img.shape[1], 1), order=_scipy_order[interpolation])
img = cv2.resize (img, (W_lr,H_lr), 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) )
self._img = img
@ -223,18 +184,14 @@ class ImageProcessor:
dtype = self.get_dtype()
self.to_ufloat32()
xp, sp, img = self._xp, self._sp, self._img
img = self._img
N,H,W,C = img.shape
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
if xp == np:
img_blur = cv2.medianBlur(img, size)
img = ne.evaluate('img*(1.0-power) + img_blur*power')
else:
img_blur = sp.ndimage.median_filter(img, size=(size,size,1) )
img = img*(1.0-power) + img_blur*power
img_blur = cv2.medianBlur(img, size)
img = ne.evaluate('img*(1.0-power) + img_blur*power')
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
self._img = img
@ -250,32 +207,23 @@ class ImageProcessor:
fade_to_border(False) clip the image in order
to fade smoothly to the border with specified blur amount
"""
xp, sp = self._xp, self._sp
erode, blur = int(erode), int(blur)
img = self._img
dtype = img.dtype
N,H,W,C = img.shape
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
img = xp.pad (img, ( (H,H), (W,W), (0,0) ) )
img = np.pad (img, ( (H,H), (W,W), (0,0) ) )
if erode > 0:
el = xp.asarray(cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3)))
el = np.asarray(cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3)))
iterations = max(1,erode//2)
if self._xp == np:
img = cv2.erode(img, el, iterations = iterations )
else:
img = sp.ndimage.binary_erosion(img, el[...,None], iterations = iterations, brute_force=True ).astype(dtype)
img = cv2.erode(img, el, iterations = iterations )
elif erode < 0:
el = xp.asarray(cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3)))
el = np.asarray(cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3)))
iterations = max(1,-erode//2)
if self._xp == np:
img = cv2.dilate(img, el, iterations = iterations )
else:
img = sp.ndimage.binary_dilation(img, el[...,None], iterations = iterations, brute_force=True).astype(dtype)
img = cv2.dilate(img, el, iterations = iterations )
if fade_to_border:
h_clip_size = H + blur // 2
@ -287,13 +235,8 @@ class ImageProcessor:
if blur > 0:
sigma = blur * 0.125 * 2
if self._xp == np:
img = cv2.GaussianBlur(img, (0, 0), sigma)
else:
img = sp.ndimage.gaussian_filter(img, (sigma, sigma,0), mode='constant')
img = cv2.GaussianBlur(img, (0, 0), sigma)
#if img.ndim == 2:
# img = img[...,None]
img = img[H:-H,W:-W]
img = img.reshape( (H,W,N,C) ).transpose( (2,0,1,3) )
@ -301,15 +244,15 @@ class ImageProcessor:
return self
def rotate90(self) -> 'ImageProcessor':
self._img = self._xp.rot90(self._img, k=1, axes=(1,2) )
self._img = np.rot90(self._img, k=1, axes=(1,2) )
return self
def rotate180(self) -> 'ImageProcessor':
self._img = self._xp.rot90(self._img, k=2, axes=(1,2) )
self._img = np.rot90(self._img, k=2, axes=(1,2) )
return self
def rotate270(self) -> 'ImageProcessor':
self._img = self._xp.rot90(self._img, k=3, axes=(1,2) )
self._img = np.rot90(self._img, k=3, axes=(1,2) )
return self
def flip_horizontal(self) -> 'ImageProcessor':
@ -330,11 +273,7 @@ class ImageProcessor:
"""
"""
xp = self._xp
img = self._img
img = xp.pad(img, ( (0,0), (t_h,b_h), (l_w,r_w), (0,0) ))
self._img = img
self._img = np.pad(self._img, ( (0,0), (t_h,b_h), (l_w,r_w), (0,0) ))
return self
def pad_to_next_divisor(self, dw=None, dh=None) -> 'ImageProcessor':
@ -343,7 +282,6 @@ class ImageProcessor:
dw,dh int
"""
xp = self._xp
img = self._img
_,H,W,_ = img.shape
@ -360,24 +298,18 @@ class ImageProcessor:
h_pad = dh - h_pad
if w_pad != 0 or h_pad != 0:
img = xp.pad(img, ( (0,0), (0,h_pad), (0,w_pad), (0,0) ))
img = np.pad(img, ( (0,0), (0,h_pad), (0,w_pad), (0,0) ))
self._img = img
return self
def sharpen(self, factor : float, kernel_size=3) -> 'ImageProcessor':
xp = self._xp
img = self._img
N,H,W,C = img.shape
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
if xp == np:
blur = cv2.GaussianBlur(img, (kernel_size, kernel_size) , 0)
img = cv2.addWeighted(img, 1.0 + (0.5 * factor), blur, -(0.5 * factor), 0)
else:
raise
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
@ -394,8 +326,6 @@ class ImageProcessor:
zero dim will be set to 1
"""
xp = self._xp
format = format.upper()
img = self._img
@ -418,7 +348,7 @@ class ImageProcessor:
transpose_order = [ d[s] for s in format ]
img = img.transpose(transpose_order)
return xp.ascontiguousarray(img)
return np.ascontiguousarray(img)
def ch(self, TC : int) -> 'ImageProcessor':
"""
@ -426,7 +356,6 @@ class ImageProcessor:
TC int >= 1
"""
xp = self._xp
img = self._img
N,H,W,C = img.shape
@ -436,7 +365,7 @@ class ImageProcessor:
if TC > C:
# Ch expand
img = img[...,0:1] # Clip to single ch first.
img = xp.repeat (img, TC, -1) # Expand by repeat
img = np.repeat (img, TC, -1) # Expand by repeat
elif TC < C:
# Ch reduction clip
img = img[...,:TC]
@ -448,7 +377,7 @@ class ImageProcessor:
"""
Converts 3 ch bgr to grayscale.
"""
img, xp = self._img, self._xp
img = self._img
_,_,_,C = img.shape
if C != 1:
dtype = self.get_dtype()
@ -458,7 +387,7 @@ class ImageProcessor:
elif C >= 3:
img = img[...,:3]
img = xp.dot(img, xp.array([0.1140, 0.5870, 0.2989], xp.float32)) [...,None]
img = np.dot(img, np.array([0.1140, 0.5870, 0.2989], np.float32)) [...,None]
img = img.astype(dtype)
self._img = img
@ -468,8 +397,6 @@ class ImageProcessor:
"""
resize to (W,H)
"""
xp, sp = self._xp, self._sp
img = self._img
N,H,W,C = img.shape
@ -479,12 +406,7 @@ class ImageProcessor:
interpolation = ImageProcessor.Interpolation.LINEAR
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
if self._xp == np:
img = cv2.resize (img, (TW, TH), interpolation=_cv_inter[interpolation])
else:
img = sp.ndimage.zoom(img, (TW/W, TH/H, 1), order=_scipy_order[interpolation])
img = cv2.resize (img, (TW, TH), interpolation=_cv_inter[interpolation])
img = img.reshape( (TH,TW,N,C) ).transpose( (2,0,1,3) )
if new_ip:
@ -498,26 +420,15 @@ class ImageProcessor:
"""
img HWC
"""
xp, sp, img = self._xp, self._sp, self._img
img = self._img
N,H,W,C = img.shape
img = img.transpose( (1,2,0,3) ).reshape( (H,W,N*C) )
if interpolation is None:
interpolation = ImageProcessor.Interpolation.LINEAR
if xp == np:
img = cv2.warpAffine(img, mat, (out_width, out_height), flags=_cv_inter[interpolation] )
else:
# AffineMat inverse
xp_mat = xp.get_array_module(mat)
mat = xp_mat.linalg.inv(xp_mat.concatenate( ( mat, xp_mat.array([[0,0,1]], xp_mat.float32)), 0) )[0:2,:]
img = cv2.warpAffine(img, mat, (out_width, out_height), flags=_cv_inter[interpolation] )
mx, my = xp.meshgrid( xp.arange(0, out_width, dtype=xp.float32), xp.arange(0, out_height, dtype=xp.float32) )
coords = xp.concatenate( (mx[None,...], my[None,...], xp.ones( (1, out_height,out_width), dtype=xp.float32)), 0 )
mat_coords = xp.matmul (xp.asarray(mat), coords.reshape( (3,-1) ) ).reshape( (2,out_height,out_width))
img = xp.concatenate([sp.ndimage.map_coordinates( img[...,c], mat_coords[::-1,...], order=_scipy_order[interpolation], mode='opencv' )[...,None] for c in range(N*C) ], -1)
img = img.reshape( (out_height,out_width,N,C) ).transpose( (2,0,1,3) )
self._img = img
return self
@ -531,23 +442,20 @@ class ImageProcessor:
"""
change image format to float32
"""
xp = self._xp
self._img = self._img.astype(xp.float32)
self._img = self._img.astype(np.float32)
return self
def as_uint8(self) -> 'ImageProcessor':
"""
change image format to uint8
"""
xp = self._xp
self._img = self._img.astype(xp.uint8)
self._img = self._img.astype(np.uint8)
return self
def to_dtype(self, dtype) -> 'ImageProcessor':
xp = self._xp
if dtype == xp.float32:
if dtype == np.float32:
return self.to_ufloat32()
elif dtype == xp.uint8:
elif dtype == np.uint8:
return self.to_uint8()
else:
raise ValueError('unsupported dtype')
@ -558,9 +466,8 @@ class ImageProcessor:
if current image dtype uint8, then image will be divided by / 255.0
otherwise no operation
"""
xp = self._xp
if self._img.dtype == xp.uint8:
self._img = self._img.astype(xp.float32)
if self._img.dtype == np.uint8:
self._img = self._img.astype(np.float32)
self._img /= 255.0
return self
@ -571,17 +478,13 @@ class ImageProcessor:
if current image dtype is float32/64, then image will be multiplied by *255
"""
xp = self._xp
img = self._img
if img.dtype in [xp.float32, xp.float64]:
if img.dtype in [np.float32, np.float64]:
img *= 255.0
img[img < 0] = 0
img[img > 255] = 255
img = img.astype(xp.uint8, copy=False)
self._img = img
np.clip(img, 0, 255, out=img)
self._img = img.astype(np.uint8, copy=False)
return self
class Interpolation(IntEnum):
@ -589,7 +492,4 @@ class ImageProcessor:
CUBIC = 1
_cv_inter = { ImageProcessor.Interpolation.LINEAR : cv2.INTER_LINEAR,
ImageProcessor.Interpolation.CUBIC : cv2.INTER_CUBIC }
_scipy_order = { ImageProcessor.Interpolation.LINEAR : 1,
ImageProcessor.Interpolation.CUBIC : 3 }
ImageProcessor.Interpolation.CUBIC : cv2.INTER_CUBIC }