Detangutify / features.py
raycosine
new augmentation
f158b5d
import numpy as np
from skimage.filters import threshold_otsu
from skimage.morphology import remove_small_objects
from skimage.feature import hog
from skimage.measure import moments_hu, label
from skimage.morphology import skeletonize, binary_dilation, disk
from scipy.ndimage import convolve
from skimage.morphology import binary_opening
def _prune_spurs(skel: np.ndarray, iters: int = 2) -> np.ndarray:
"""
迭代剪掉骨架上长度很短的端点分支(spur)。
iters 表示最多向内剪掉几步(像素)。推荐 1~3。
"""
s = skel.copy().astype(bool)
# 用 3x3 邻域统计端点:中心权重10,其它1;“10+1=11”即1个邻居的端点
K = np.array([[1,1,1],
[1,10,1],
[1,1,1]], dtype=np.uint8)
for _ in range(iters):
nb = convolve(s.astype(np.uint8), K, mode="constant", cval=0)
endpoints = (nb == 11) # 只有 1 个邻居
# 只剪 endpoints,不动分叉/主干
s = s & ~endpoints
return s
def _ensure_ink_true(bw_bool: np.ndarray) -> np.ndarray:
bw = bw_bool.astype(bool)
if bw.mean() > 0.5:
bw = ~bw
return bw
def stroke_normalize(bw: np.ndarray, target_px: int = 2) -> np.ndarray:
if bw.dtype != bool:
bw = (bw > 0)
if bw.mean() > 0.5:
bw = ~bw
skel = skeletonize(bw)
skel = _prune_spurs(skel, iters=2) # ← 新增:剪短刺,去笔锋小尖
if target_px <= 1:
return skel.astype(np.float32)
rad = max(1, int(round(target_px/2)))
thick = binary_dilation(skel, disk(rad))
#thick = binary_opening(thick, disk(1))#optional
return (thick & bw).astype(np.float32)
def to_64_gray(imgPIL):
return np.array(imgPIL, dtype=np.uint8)
def binarize(gray64, min_size: int = 4, keep_largest: bool = False):
if gray64.ndim == 3:
gray64 = gray64[..., 0]
g = gray64.astype(float)
if g.max() > 1:
g /= 255.0
t = threshold_otsu(g)
bw = (g <= t)
bw = remove_small_objects(bw.astype(bool), min_size=min_size).astype(bool)
if keep_largest:
lab = label(bw)
if lab.max() > 0:
areas = np.bincount(lab.ravel()); areas[0] = 0
bw = (lab == areas.argmax())
return bw
def proj_features(bw):
# normalization
hp = bw.sum(axis=0); vp = bw.sum(axis=1)
if hp.max()>0: hp = hp/hp.max()
if vp.max()>0: vp = vp/vp.max()
# fix dimension
def pool(v, m=32):
idx = np.linspace(0, len(v), m+1, endpoint=True).astype(int)
return np.array([v[idx[i]:idx[i+1]].mean() for i in range(m)], dtype=np.float32)
return np.concatenate([pool(hp), pool(vp)])
def feat_vec(bw):
hu = moments_hu(bw).astype(np.float32)
hu = np.sign(hu)*np.log1p(np.abs(hu))
hogv = hog(bw, orientations=9, pixels_per_cell=(8,8), cells_per_block=(2,2),
block_norm='L2-Hys', feature_vector=True).astype(np.float32)
proj = proj_features(bw).astype(np.float32)
v = np.concatenate([hu, hogv, proj]).astype(np.float32)
n = np.linalg.norm(v) + 1e-8
return v / n
def cosine_sim(a, B):
return (B @ a) / (np.linalg.norm(a)+1e-8) / (np.linalg.norm(B,axis=1)+1e-8)