Spaces:
Sleeping
Sleeping
File size: 3,277 Bytes
f24d19f 1213d9c df01e43 f24d19f 1213d9c f24d19f 1213d9c f24d19f 1213d9c 7c0ca8f f24d19f 316424f f24d19f 1213d9c f24d19f 3272647 f24d19f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import os
import sys
import cv2
import numpy as np
from PIL import Image
debug = False
def gray3d_to_2d(grayscale: np.ndarray) -> np.ndarray:
channel = grayscale.shape[2] if grayscale.ndim == 3 else 1
if channel!=1:
text = f"grayscale shape = {grayscale.shape} channel = {channel} ndim = {grayscale.ndim} size = {grayscale.size}"
raise ValueError(f"color maybe rgb or rgba {text}")
if grayscale.ndim == 2:
return grayscale
return np.squeeze(grayscale)
def pil_to_cv(image):
cv_image = np.array(image, dtype=np.uint8)
if cv_image.shape[2] == 3: # カラー
cv_image = cv2.cvtColor(cv_image, cv2.COLOR_RGB2BGR)
elif cv_image.shape[2] == 4: #
cv_image = cv2.cvtColor(cv_image, cv2.COLOR_RGBA2BGR)
return cv_image
def blend_rgb_images(image1: np.ndarray, image2: np.ndarray, mask: np.ndarray) -> np.ndarray:
if image1.shape != image2.shape or image1.shape[:2] != mask.shape:
raise ValueError("not same shape")
# 画像を float 型に変換
image1 = image1.astype(float)
image2 = image2.astype(float)
# mask to 3 chan 0 -1 value
alpha = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR).astype(float) / 255.0
# calcurate blend
blended = (1 - alpha) * image1 + alpha * image2
return blended.astype(np.uint8)
def process_cvinpaint(image,mask_image,inpaint_radius,blur_radius,edge_expand,inpaint_mode,dilate=0):
#print("process cvinpaint")
#print(blur_radius,",",edge_expand)
cv_image = pil_to_cv(image)
cv_mask = pil_to_cv(mask_image)
cv_gray = cv2.cvtColor(cv_mask,cv2.COLOR_BGR2GRAY)
mask = gray3d_to_2d(cv_gray)
if dilate>0:
kernel = np.ones((dilate, dilate), np.uint8)
mask = cv2.dilate(mask, kernel, iterations=1)
#cv2.imwrite("_mask.jpg",mask)
#cv2.imwrite("_image.jpg",cv_image)
mode = cv2.INPAINT_TELEA if inpaint_mode == "Telea" else cv2.INPAINT_NS
img_inpainted = cv2.inpaint(cv_image, mask,inpaint_radius, mode)
if debug:
cv2.imwrite("close_eye_inpaint.jpg",img_inpainted)
## blur
if blur_radius > 0:
if blur_radius%2==0:
blur_radius += 1
#print(blur_radius)
blurred_image = cv2.GaussianBlur(img_inpainted, (blur_radius, blur_radius), 0) #should be odd
if debug:
cv2.imwrite("close_eye_inpaint_burred.jpg",blurred_image)
else:
blurred_image = img_inpainted
# expand edge and blur
kernel = np.ones((edge_expand, edge_expand), np.uint8)
extend_mask = cv2.dilate(mask, kernel, iterations=1)
if edge_expand > 0 and blur_radius > 0:
extend_burred_mask = cv2.GaussianBlur(extend_mask, (blur_radius, blur_radius), 0)
else:
extend_burred_mask = extend_mask
img_inpainted = blend_rgb_images(img_inpainted,blurred_image,extend_burred_mask)
output_image = img_inpainted.copy()
if output_image.shape[2] == 3: # カラー
output_image = cv2.cvtColor(output_image, cv2.COLOR_BGR2RGB)
return Image.fromarray(output_image),Image.fromarray(mask)
if __name__ == "__main__":
image = Image.open(sys.argv[1])
mask = Image.open(sys.argv[2])
output = process_cvinpaint(image,mask)
output.save(sys.argv[3]) |