IOPaint/iopaint/model/helper/g_diffuser_bot.py

63 lines
1.9 KiB
Python
Raw Normal View History

2023-08-30 07:28:31 +02:00
import cv2
import numpy as np
2024-04-24 14:22:29 +02:00
def expand_image(cv2_img, top: int, right: int, bottom: int, left: int):
2023-11-15 10:18:32 +01:00
assert cv2_img.shape[2] == 3
2023-08-30 07:28:31 +02:00
origin_h, origin_w = cv2_img.shape[:2]
2023-11-15 10:18:32 +01:00
# TODO: which is better?
2024-04-24 14:22:29 +02:00
# new_img = np.ones((new_height, new_width, 3), np.uint8) * 255
2023-11-15 10:18:32 +01:00
new_img = cv2.copyMakeBorder(
cv2_img, top, bottom, left, right, cv2.BORDER_REPLICATE
2023-08-30 07:28:31 +02:00
)
2024-04-24 14:22:29 +02:00
inner_padding_left = 0 if left > 0 else 0
inner_padding_right = 0 if right > 0 else 0
inner_padding_top = 0 if top > 0 else 0
inner_padding_bottom = 0 if bottom > 0 else 0
2024-03-05 15:23:53 +01:00
mask_image = np.zeros(
(
2024-04-24 14:22:29 +02:00
origin_h - inner_padding_top - inner_padding_bottom,
origin_w - inner_padding_left - inner_padding_right,
2024-03-05 15:23:53 +01:00
),
2024-04-24 14:22:29 +02:00
np.uint8,
)
2024-03-05 15:23:53 +01:00
mask_image = cv2.copyMakeBorder(
mask_image,
top + inner_padding_top,
bottom + inner_padding_bottom,
left + inner_padding_left,
right + inner_padding_right,
cv2.BORDER_CONSTANT,
2024-04-24 14:22:29 +02:00
value=255,
2024-03-05 15:23:53 +01:00
)
# k = 2*int(min(origin_h, origin_w) // 6)+1
2024-04-24 14:22:29 +02:00
# k = 7
# mask_image = cv2.GaussianBlur(mask_image, (k, k), 0)
2024-03-05 15:23:53 +01:00
return new_img, mask_image
2023-08-30 07:28:31 +02:00
if __name__ == "__main__":
from pathlib import Path
current_dir = Path(__file__).parent.absolute().resolve()
2024-03-05 15:23:53 +01:00
image_path = "/Users/cwq/code/github/IOPaint/iopaint/tests/bunny.jpeg"
2023-08-30 07:28:31 +02:00
init_image = cv2.imread(str(image_path))
2024-04-24 14:22:29 +02:00
init_image, mask_image = expand_image(
2023-08-30 07:28:31 +02:00
init_image,
2024-03-05 15:23:53 +01:00
top=0,
right=0,
bottom=0,
2023-11-14 07:19:56 +01:00
left=100,
softness=20,
space=20,
2023-08-30 07:28:31 +02:00
)
print(mask_image.dtype, mask_image.min(), mask_image.max())
print(init_image.dtype, init_image.min(), init_image.max())
mask_image = mask_image.astype(np.uint8)
init_image = init_image.astype(np.uint8)
cv2.imwrite("expanded_image.png", init_image)
cv2.imwrite("expanded_mask.png", mask_image)