diff --git a/main.py b/main.py index b3f9c04..3ad7d2a 100755 --- a/main.py +++ b/main.py @@ -77,7 +77,7 @@ def main(only_spec=False, only_color=False): mask_spec = spec_detector.predict(img_data) mask_rgb = rgb_detector.predict(rgb_data) # 进行喷阀的合并 - masks = [utils.valve_merge(mask, merge_size=Config.valve_merge_size) for mask in [mask_spec, mask_rgb]] + masks = [utils.valve_expend(mask) for mask in [mask_spec, mask_rgb]] # control the size of the output masks, 在resize前,图像的宽度是和喷阀对应的 masks = [cv2.resize(mask.astype(np.uint8), Config.target_size) for mask in masks] # 写出 diff --git a/utils.py b/utils.py index 538082e..57183cb 100755 --- a/utils.py +++ b/utils.py @@ -153,13 +153,18 @@ def size_threshold(img, blk_size, threshold, last_end: np.ndarray = None) -> np. def valve_merge(img: np.ndarray, merge_size: int = 2) -> np.ndarray: assert img.shape[1] % merge_size == 0 # 列数必须能够被整除 img_shape = (img.shape[1], img.shape[0]) - img = img.reshape((img.shape[0], img.shape[1]//merge_size, merge_size)) + img = img.reshape((img.shape[0], img.shape[1] // merge_size, merge_size)) img = np.sum(img, axis=2) img[img > 0] = 1 img = cv2.resize(img.astype(np.uint8), dsize=img_shape) return img +def valve_expend(img: np.ndarray) -> np.ndarray: + kernel = np.ones((1, 3), np.uint8) + return cv2.dilate(img, kernel) + + def read_envi_ascii(file_name, save_xy=False, hdr_file_name=None): """ Read envi ascii file. Use ENVI ROI Tool -> File -> output ROIs to ASCII...