From 597284c3e1ef3f9062dee5600fd9254a713ce838 Mon Sep 17 00:00:00 2001 From: "li.zhenye" <李> Date: Wed, 27 Jul 2022 13:40:35 +0800 Subject: [PATCH] remove the saving function --- main.py | 107 +------------------------------------------------------- 1 file changed, 1 insertion(+), 106 deletions(-) diff --git a/main.py b/main.py index 32602f9..7ac3b21 100755 --- a/main.py +++ b/main.py @@ -8,9 +8,6 @@ from models import RgbDetector, SpecDetector, ManualTree, AnonymousColorDetector import cv2 -SAVE_IMG, SAVE_NUM = False, 30 - - def main(): spec_detector = SpecDetector(blk_model_path=Config.blk_model_path, pixel_model_path=Config.pixel_model_path) rgb_detector = RgbDetector(tobacco_model_path=Config.rgb_tobacco_model_path, @@ -23,8 +20,6 @@ def main(): os.mkfifo(mask_fifo_path, 0o777) if not os.access(rgb_fifo_path, os.F_OK): os.mkfifo(rgb_fifo_path, 0o777) - if SAVE_IMG: - img_list = [] while True: fd_img = os.open(img_fifo_path, os.O_RDONLY) fd_rgb = os.open(rgb_fifo_path, os.O_RDONLY) @@ -55,11 +50,6 @@ def main(): img_data = np.frombuffer(data_total, dtype=np.float32).reshape((Config.nRows, Config.nBands, -1)) \ .transpose(0, 2, 1) rgb_data = np.frombuffer(rgb_data_total, dtype=np.uint8).reshape((Config.nRgbRows, Config.nRgbCols, -1)) - if SAVE_IMG: - SAVE_NUM -= 1 - img_list.append((rgb_data, img_data)) - if SAVE_NUM <= 0: - break # 光谱识别 mask = spec_detector.predict(img_data) # rgb识别 @@ -71,108 +61,13 @@ def main(): # mask_result = mask_rgb.astype(np.uint8) mask_result = mask_result.repeat(Config.blk_size, axis=0).repeat(Config.blk_size, axis=1).astype(np.uint8) t2 = time.time() - print(f'rgb len = {len(rgb_data)}') # 写出 fd_mask = os.open(mask_fifo_path, os.O_WRONLY) os.write(fd_mask, mask_result.tobytes()) os.close(fd_mask) t3 = time.time() - print(f'total time is:{t3 - t1}') - for i, img in enumerate(img_list): - print(f"writing img {i}...") - cv2.imwrite(f"./{i}.png", img[0][..., ::-1]) - np.save(f'./{i}.npy', img[1]) - i += 1 - - - -def save_main(): - threshold = Config.spec_size_threshold - rgb_threshold = Config.rgb_size_threshold - manual_tree = ManualTree(blk_model_path=Config.blk_model_path, pixel_model_path=Config.pixel_model_path) - tobacco_detector = AnonymousColorDetector(file_path=Config.rgb_tobacco_model_path) - background_detector = AnonymousColorDetector(file_path=Config.rgb_background_model_path) - total_len = Config.nRows * Config.nCols * Config.nBands * 4 # float型变量, 4个字节 - total_rgb = Config.nRgbRows * Config.nRgbCols * Config.nRgbBands * 1 # int型变量 - if not os.access(img_fifo_path, os.F_OK): - os.mkfifo(img_fifo_path, 0o777) - if not os.access(mask_fifo_path, os.F_OK): - os.mkfifo(mask_fifo_path, 0o777) - if not os.access(rgb_fifo_path, os.F_OK): - os.mkfifo(rgb_fifo_path, 0o777) - img_list = [] - idx = 0 - while idx <= 30: - idx += 1 - fd_img = os.open(img_fifo_path, os.O_RDONLY) - fd_rgb = os.open(rgb_fifo_path, os.O_RDONLY) - data = os.read(fd_img, total_len) - - # 读取(开启一个管道) - if len(data) < 3: - threshold = int(float(data)) - print("[INFO] Get threshold: ", threshold) - continue - else: - data_total = data - rgb_data = os.read(fd_rgb, total_rgb) - if len(rgb_data) < 3: - rgb_threshold = int(float(rgb_data)) - print(rgb_threshold) - continue - else: - rgb_data_total = rgb_data - os.close(fd_img) - os.close(fd_rgb) - - # 识别 - t1 = time.time() - img_data = np.frombuffer(data_total, dtype=np.float32).reshape((Config.nRows, Config.nBands, -1)). \ - transpose(0, 2, 1) - rgb_data = np.frombuffer(rgb_data_total, dtype=np.uint8).reshape((Config.nRgbRows, Config.nRgbCols, -1)) - img_list.append((rgb_data.copy(), img_data.copy())) - - pixel_predict_result = manual_tree.pixel_predict_ml_dilation(data=img_data, iteration=1) - blk_predict_result = manual_tree.blk_predict(data=img_data) - rgb_data = tobacco_detector.pretreatment(rgb_data) - # print(rgb_data.shape) - rgb_predict_result = 1 - (background_detector.predict(rgb_data, threshold_low=Config.threshold_low, - threshold_high=Config.threshold_high) | - tobacco_detector.swell(tobacco_detector.predict(rgb_data, - threshold_low=Config.threshold_low, - threshold_high=Config.threshold_high))) - mask_rgb = rgb_predict_result.reshape(Config.nRows, Config.nCols // Config.blk_size, Config.blk_size) \ - .sum(axis=2).reshape(Config.nRows // 4, Config.blk_size, Config.nCols // Config.blk_size) \ - .sum(axis=1) - mask_rgb[mask_rgb <= rgb_threshold] = 0 - mask_rgb[mask_rgb > rgb_threshold] = 1 - mask = (pixel_predict_result & blk_predict_result).astype(np.uint8) - mask = mask.reshape(Config.nRows, Config.nCols // Config.blk_size, Config.blk_size) \ - .sum(axis=2).reshape(Config.nRows // 4, Config.blk_size, Config.nCols // Config.blk_size) \ - .sum(axis=1) - mask[mask <= threshold] = 0 - mask[mask > threshold] = 1 - mask_result = (mask | mask_rgb).astype(np.uint8) - # mask_result = mask_rgb - mask_result = mask_result.repeat(Config.blk_size, axis=0).repeat(Config.blk_size, axis=1).astype(np.uint8) - t2 = time.time() - print(f'rgb len = {len(rgb_data)}') - - # 写出 - fd_mask = os.open(mask_fifo_path, os.O_WRONLY) - os.write(fd_mask, mask_result.tobytes()) - os.close(fd_mask) - t3 = time.time() - print(f'total time is:{t3 - t1}') - i = 0 - print("Stop Serving") - for img in img_list: - print(f"writing img {i}...") - cv2.imwrite(f"./{i}.png", img[0][..., ::-1]) - np.save(f'./{i}.npy', img[1]) - i += 1 - print("save success") + print(f'total time is:{t3 - t1}\n') if __name__ == '__main__':