import cv2, numpy as np
from skimage.morphology import skeletonize

src = "/home/gxs/Desktop/data/2300图像文件测试数据/R4407845900002021120182/R4407845900002021120182_01_01_01_01.bmp"
img = cv2.imread(src, cv2.IMREAD_GRAYSCALE)

# upscale 2x
big = cv2.resize(img, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)
# strong non-local-means denoise
den = cv2.fastNlMeansDenoising(big, None, h=25, templateWindowSize=7, searchWindowSize=21)

# adaptive threshold on denoised 2x image (ridge period ~16px now)
a = cv2.adaptiveThreshold(den, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                          cv2.THRESH_BINARY, 81, 3)
ridge = (a == 0).astype(np.uint8)

# print region from global otsu
_, gotsu = cv2.threshold(cv2.GaussianBlur(den, (5,5), 0), 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
pr = (gotsu == 0).astype(np.uint8)
if pr.mean() > 0.5:
    pr = 1 - pr
ridge = ridge * pr

k2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2,2))
k3 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
ridge = cv2.morphologyEx(ridge, cv2.MORPH_OPEN, k2)
ridge = cv2.morphologyEx(ridge, cv2.MORPH_CLOSE, k3, iterations=2)

# keep big components
n, labels, stats, _ = cv2.connectedComponentsWithStats(ridge, 8)
keep = np.zeros_like(ridge)
for i in range(1, n):
    if stats[i, cv2.CC_STAT_AREA] >= 400:
        keep[labels == i] = 1
ridge = keep
print(f"ridge fraction: {ridge.mean():.3f}")
skel = skeletonize(ridge > 0)
print(f"skeleton px: {skel.sum()}")
canvas = np.hstack([cv2.resize(img, (640, 640)),
                    cv2.resize((ridge*255).astype(np.uint8), (640, 640)),
                    cv2.resize(((skel>0)*255).astype(np.uint8), (640, 640))])
cv2.imwrite('/tmp/fp_nlm.png', canvas)
print('saved')
