raycv/tool/img2skeleton.py

42 lines
1.1 KiB
Python

from pathlib import Path
import fire
import cv2
import numpy as np
from tqdm import tqdm
def xDoG(img, sigma, k_sigma, p, epsilon, phi):
sigma_large = sigma * k_sigma
sharped_img = (1 + p) * cv2.GaussianBlur(img, (0, 0), sigma) - p * cv2.GaussianBlur(img, (0, 0), sigma_large)
img = np.multiply(img, sharped_img)
t = np.zeros(img.shape)
t[img >= epsilon] = 1.0
img_dark_indices = img < epsilon
t[img_dark_indices] = 1.0 + np.tanh(phi * (img[img_dark_indices] - epsilon))
return t * 256
def transform_single(origin, to, anime: bool = True):
img = cv2.imread(str(origin), cv2.IMREAD_GRAYSCALE)
if anime:
r = xDoG(img / 256, 0.3, 4.5, 19, 0.01, 10 ^ 8 * 2)
else:
r = xDoG(img / 256, 0.7, 5, 20, 0.02, 10 ^ 8 * 2)
cv2.imwrite(str(to), r)
def transform(origin, to, anime=True):
origin = Path(origin)
to = Path(to)
if origin.is_dir() and to.is_dir():
for f in tqdm(origin.glob("*")):
transform_single(f, to / f.name, anime)
elif origin.is_file():
transform_single(origin, to, anime)
if __name__ == '__main__':
fire.Fire(transform)