diff --git a/config.py b/config.py index 7a318a3..7685ea9 100644 --- a/config.py +++ b/config.py @@ -1,4 +1,5 @@ -import torch +import os + import numpy as np @@ -36,3 +37,6 @@ class Config: # save part offset_vertical = 0 + + # logging + root_dir = os.path.split(os.path.realpath(__file__))[0] diff --git a/main.py b/main.py index 17314d8..b3f9c04 100755 --- a/main.py +++ b/main.py @@ -1,4 +1,6 @@ import os +import sys + import cv2 import time import numpy as np @@ -6,6 +8,7 @@ import numpy as np import utils from config import Config from models import RgbDetector, SpecDetector +import logging def main(only_spec=False, only_color=False): @@ -34,27 +37,34 @@ def main(only_spec=False, only_color=False): if len(data) < 3: threshold = int(float(data)) Config.spec_size_threshold = threshold - print("[INFO] Get spec threshold: ", threshold) + logging.info('[INFO] Get spec threshold: ', threshold) else: data_total = data os.close(fd_img) - # rgb data read rgb_data = os.read(fd_rgb, total_rgb) if len(rgb_data) < 3: - rgb_threshold = int(float(rgb_data)) - Config.rgb_size_threshold = rgb_threshold - print("[INFO] Get rgb threshold", rgb_threshold) + try: + rgb_threshold = int(float(rgb_data)) + Config.rgb_size_threshold = rgb_threshold + logging.info(f'Get rgb threshold: {rgb_threshold}') + except Exception as e: + logging.error(f'毁灭性错误:收到长度小于3却无法转化为整数的数据,{e}.') continue else: rgb_data_total = rgb_data os.close(fd_rgb) - # 识别 - t1 = time.time() - img_data = np.frombuffer(data_total, dtype=np.float32).reshape((Config.nRows, Config.nBands, -1)) \ + since = time.time() + try: + 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)) + except Exception as e: + logging.error(f'毁灭性错误!收到的光谱数据长度为{len(data_total)}无法转化成指定的形状 {e}') + try: + rgb_data = np.frombuffer(rgb_data_total, dtype=np.uint8).reshape((Config.nRgbRows, Config.nRgbCols, -1)) + except Exception as e: + logging.error(f'毁灭性错误!收到的rgb数据长度为{len(rgb_data)}无法转化成指定形状 {e}') if only_spec: # 光谱识别 mask_spec = spec_detector.predict(img_data) @@ -66,7 +76,6 @@ def main(only_spec=False, only_color=False): else: 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]] # control the size of the output masks, 在resize前,图像的宽度是和喷阀对应的 @@ -77,8 +86,10 @@ def main(only_spec=False, only_color=False): fd_mask = os.open(fifo, os.O_WRONLY) os.write(fd_mask, mask.tobytes()) os.close(fd_mask) - t3 = time.time() - print(f'total time is:{t3 - t1}') + time_spent = (time.time() - since) * 1000 + logging.info(f'Total time is: {time_spent:.2f} ms') + if time_spent > 200: + logging.warning(f'警告预测超时,预测耗时超过了200ms,The prediction time is {time_spent:.2f} ms.') if __name__ == '__main__': @@ -86,11 +97,22 @@ if __name__ == '__main__': parser = argparse.ArgumentParser(description='主程序') parser.add_argument('-oc', default=False, action='store_true', help='只进行RGB彩色预测 only rgb', required=False) parser.add_argument('-os', default=False, action='store_true', help='只进行光谱预测 only spec', required=False) + parser.add_argument('-d', default=False, action='store_true', help='是否使用DEBUG模式', required=False) args = parser.parse_args() # fifo 参数 - img_fifo_path = "/tmp/dkimg.fifo" - rgb_fifo_path = "/tmp/dkrgb.fifo" + img_fifo_path = '/tmp/dkimg.fifo' + rgb_fifo_path = '/tmp/dkrgb.fifo' # mask fifo - mask_fifo_path = "/tmp/dkmask.fifo" - rgb_mask_fifo_path = "/tmp/dkmask_rgb.fifo" + mask_fifo_path = '/tmp/dkmask.fifo' + rgb_mask_fifo_path = '/tmp/dkmask_rgb.fifo' + # logging相关 + file_handler = logging.FileHandler(os.path.join(Config.root_dir, '.tobacco_algorithm.log')) + file_handler.setLevel(logging.WARNING) + console_handler = logging.StreamHandler(sys.stdout) + if args.d: + console_handler.setLevel(logging.DEBUG) + else: + console_handler.setLevel(logging.WARNING) + logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', + handlers=[file_handler, console_handler]) main(only_spec=args.os, only_color=args.oc)