Spaces:
Sleeping
Sleeping
| import traceback | |
| import numpy as np | |
| import cv2 as cv | |
| from matplotlib import pyplot as plt | |
| def get_binary_img_(img): | |
| gray_img = img | |
| if len(img.shape) > 2: | |
| gray_img = cv.cvtColor(img, cv.COLOR_BGR2GRAY) | |
| binary_img = cv.Canny(gray_img, 80, 150) | |
| return binary_img | |
| def get_morp_dilate_(binary_img): | |
| kernel = cv.getStructuringElement(cv.MORPH_RECT, (3, 3)) | |
| # morp_dilate = cv.morphologyEx(binaryImg, cv.MORPH_DILATE, kernel=(1, 3), iterations=3) | |
| # morp_dilate = cv.morphologyEx(morp_dilate, cv.MORPH_DILATE, kernel=(3, 1), iterations=3) | |
| # morp_dilate = cv.morphologyEx(binaryImg, cv.MORPH_DILATE, kernel=(11, 11), iterations=3) | |
| morp_dilate = cv.morphologyEx(binary_img, cv.MORPH_DILATE, kernel=kernel, iterations=3) | |
| return morp_dilate | |
| def get_water_img_(img, morp_dilate): | |
| # 寻找图像轮廓 返回修改后的 图像的轮廓 以及它们的层次 | |
| # contours, hierarchy = cv.findContours(gray_img, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE) | |
| # contours, hierarchy = cv.findContours(thresh, cv.RETR_TREE, cv.CHAIN_APPROX_NONE) | |
| # contours, hierarchy = cv.findContours(thresh, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE) | |
| contours, hierarchy = cv.findContours(morp_dilate, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE) | |
| # 32位有符号整数类型， | |
| marks = np.zeros(morp_dilate.shape[:2], np.int32) | |
| # 绘制每一个轮廓 | |
| for index in range(len(contours)): | |
| # 对marks进行标记，对不同区域的轮廓使用不同的亮度绘制，相当于设置注水点，有多少个轮廓，就有多少个轮廓 | |
| # 图像上不同线条的灰度值是不同的，底部略暗，越往上灰度越高 | |
| marks = cv.drawContours(marks, contours, index, (index, index, index), 1, 8, hierarchy) | |
| # 使用分水岭算法 | |
| # 经过watershed函数的处理，不同区域间的值被置为-1（边界）没有标记清楚的区域被置为0，其他每个区域的值保持不变：1，2，...，contours.size() | |
| marks_water = cv.watershed(img, marks) | |
| return marks_water | |
| def get_mask_img_(morp_dilate, file_dir): | |
| contours, hierarchy = cv.findContours(morp_dilate, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE) | |
| # 32位有符号整数类型， | |
| marks = np.zeros(morp_dilate.shape[:2], np.int32) | |
| for index in range(len(contours)): | |
| dist = cv.pointPolygonTest(contours[index], (marks.shape[0] // 2, marks.shape[1] // 2), True) | |
| if dist >= 0: | |
| marks = cv.drawContours(marks, contours, contourIdx=index, color=1, thickness=1, lineType=8, | |
| hierarchy=hierarchy) | |
| edges = np.zeros((marks.shape[0] + 2, marks.shape[1] + 2), np.uint8) # 掩码，长短需要加2个像素 | |
| try: | |
| cv.floodFill(marks, edges, (marks.shape[0] // 2, marks.shape[1] // 2), 1, cv.FLOODFILL_MASK_ONLY) # 漫水填充 | |
| except Exception as e: | |
| if file_dir: | |
| print(file_dir) | |
| print(e) | |
| print("=================") | |
| print(traceback.format_exc()) | |
| # raise e | |
| marks = np.ones(morp_dilate.shape[:2], np.int32) | |
| return marks | |
| def get_binary_img(binary_img, mask): | |
| masked_binary_img = cv.bitwise_and(binary_img, binary_img, mask=mask.astype('uint8')) | |
| return masked_binary_img | |
| def get_water_img(img, morp_dilate, mask): | |
| water_img = get_water_img_(img, morp_dilate) | |
| masked_water = cv.bitwise_and(water_img, water_img, mask=mask.astype('uint8')) | |
| return masked_water | |
| def get_more_dim(img, file_dir, source_img=None): | |
| if source_img is None: | |
| source_img = img | |
| # img: ndarray: 852, 847, 3 | |
| binary_img = get_binary_img_(img) | |
| morp_dilate = get_morp_dilate_(binary_img) | |
| mask = get_mask_img_(morp_dilate, file_dir) | |
| masked_binary_img = get_binary_img(binary_img, mask) | |
| masked_water = get_water_img(source_img, morp_dilate, mask) | |
| # print(f"masked_binary_img shape:{masked_binary_img.shape} masked_water shape:{masked_water.shape}") | |
| # print(f"type(masked_binary_img):{type(masked_binary_img)} type(masked_water):{type(masked_water)}") | |
| # return np.stack((masked_binary_img, mask), axis=0) | |
| return masked_binary_img, masked_water | |