mirror of
https://github.com/NanjingForestryUniversity/supermachine--tomato-passion_fruit.git
synced 2025-11-09 14:54:07 +00:00
167 lines
6.5 KiB
Python
167 lines
6.5 KiB
Python
# -*- coding: utf-8 -*-
|
||
# @Time : 2024/4/20 18:45
|
||
# @Author : TG
|
||
# @File : main.py
|
||
# @Software: PyCharm
|
||
|
||
import sys
|
||
import os
|
||
|
||
import cv2
|
||
|
||
from root_dir import ROOT_DIR
|
||
from classifer import Spec_predict, Data_processing, ImageClassifier
|
||
import logging
|
||
from utils import Pipe
|
||
import numpy as np
|
||
import time
|
||
|
||
|
||
|
||
|
||
def process_data(cmd: str, images: list, spec: any, dp: Data_processing, pipe: Pipe, detector: Spec_predict) -> bool:
|
||
"""
|
||
处理指令
|
||
|
||
:param cmd: 指令类型
|
||
:param images: 图像数据列表
|
||
:param spec: 光谱数据
|
||
:param detector: 模型
|
||
:return: 是否处理成功
|
||
"""
|
||
diameter_axis_list = []
|
||
max_defect_num = 0 # 初始化最大缺陷数量为0
|
||
max_total_defect_area = 0 # 初始化最大总像素数为0
|
||
|
||
for i, img in enumerate(images):
|
||
if cmd == 'TO':
|
||
# 番茄
|
||
diameter, green_percentage, number_defects, total_pixels, rp = dp.analyze_tomato(img)
|
||
if i <= 2:
|
||
diameter_axis_list.append(diameter)
|
||
max_defect_num = max(max_defect_num, number_defects)
|
||
max_total_defect_area = max(max_total_defect_area, total_pixels)
|
||
if i == 1:
|
||
rp_result = rp
|
||
gp = round(green_percentage)
|
||
|
||
elif cmd == 'PF':
|
||
# 百香果
|
||
diameter, weigth, number_defects, total_pixels, rp = dp.analyze_passion_fruit(img)
|
||
if i <= 2:
|
||
diameter_axis_list.append(diameter)
|
||
max_defect_num = max(max_defect_num, number_defects)
|
||
max_total_defect_area = max(max_total_defect_area, total_pixels)
|
||
if i == 1:
|
||
rp_result = rp
|
||
weigth = weigth
|
||
|
||
else:
|
||
logging.error(f'错误指令，指令为{cmd}')
|
||
return False
|
||
|
||
diameter = round(sum(diameter_axis_list) / 3)
|
||
|
||
if cmd == 'TO':
|
||
brix = 0
|
||
weigth = 0
|
||
response = pipe.send_data(cmd=cmd, brix=brix, diameter=diameter, green_percentage=gp, weigth=weigth,
|
||
defect_num=max_defect_num, total_defect_area=max_total_defect_area, rp=rp_result)
|
||
return response
|
||
elif cmd == 'PF':
|
||
green_percentage = 0
|
||
brix = detector.predict(spec)
|
||
response = pipe.send_data(cmd=cmd, brix=brix, green_percentage=green_percentage, diameter=diameter, weigth=weigth,
|
||
defect_num=max_defect_num, total_defect_area=max_total_defect_area, rp=rp_result)
|
||
return response
|
||
|
||
def main(is_debug=False):
|
||
file_handler = logging.FileHandler(os.path.join(ROOT_DIR, 'tomato.log'), encoding='utf-8')
|
||
file_handler.setLevel(logging.DEBUG if is_debug else logging.WARNING)
|
||
console_handler = logging.StreamHandler(sys.stdout)
|
||
console_handler.setLevel(logging.DEBUG if is_debug else logging.WARNING)
|
||
logging.basicConfig(format='%(asctime)s %(filename)s[line:%(lineno)d] - %(levelname)s - %(message)s',
|
||
handlers=[file_handler, console_handler],
|
||
level=logging.DEBUG)
|
||
detector = Spec_predict(ROOT_DIR/'models'/'passion_fruit_2.joblib')
|
||
classifier = ImageClassifier(ROOT_DIR/'models'/'resnet18_0616.pth', ROOT_DIR/'models'/'class_indices.json')
|
||
dp = Data_processing()
|
||
|
||
_ = detector.predict(np.ones((30, 30, 224), dtype=np.uint16))
|
||
_ = classifier.predict(np.ones((224, 224, 3), dtype=np.uint8))
|
||
# _, _, _, _, _ =dp.analyze_tomato(cv2.imread(r'D:\project\supermachine--tomato-passion_fruit\20240529RGBtest3\data\tomato_img\bad\71.bmp'))
|
||
# _, _, _, _, _ = dp.analyze_passion_fruit(cv2.imread(r'D:\project\supermachine--tomato-passion_fruit\20240529RGBtest3\data\passion_fruit_img\38.bmp'))
|
||
print('系统初始化完成')
|
||
|
||
rgb_receive_name = r'\\.\pipe\rgb_receive'
|
||
rgb_send_name = r'\\.\pipe\rgb_send'
|
||
spec_receive_name = r'\\.\pipe\spec_receive'
|
||
pipe = Pipe(rgb_receive_name, rgb_send_name, spec_receive_name)
|
||
rgb_receive, rgb_send, spec_receive = pipe.create_pipes(rgb_receive_name, rgb_send_name, spec_receive_name)
|
||
# 预热循环，只处理cmd为'YR'的数据
|
||
while True:
|
||
start_time00 = time.time()
|
||
data = pipe.receive_rgb_data(rgb_receive)
|
||
cmd, _ = pipe.parse_img(data)
|
||
end_time00 = time.time()
|
||
print(f'接收预热数据时间：{end_time00 - start_time00}秒')
|
||
if cmd == 'YR':
|
||
break # 当接收到的不是预热命令时，结束预热循环
|
||
while True:
|
||
start_time = time.time()
|
||
images = []
|
||
cmd = None
|
||
for _ in range(5):
|
||
start_time1 = time.time()
|
||
data = pipe.receive_rgb_data(rgb_receive)
|
||
end_time10 = time.time()
|
||
print(f'接收一份数据时间：{end_time10 - start_time1}秒')
|
||
|
||
start_time11 = time.time()
|
||
cmd, img = pipe.parse_img(data)
|
||
end_time1 = time.time()
|
||
print(f'处理一份数据时间：{end_time1 - start_time11}秒')
|
||
print(f'接收一张图时间：{end_time1 - start_time1}秒')
|
||
|
||
# 使用分类器进行预测
|
||
prediction = classifier.predict(img)
|
||
print(f'预测结果：{prediction}')
|
||
if prediction == 1:
|
||
images.append(img)
|
||
else:
|
||
response = pipe.send_data(cmd='KO', brix=0, diameter=0, green_percentage=0, weigth=0, defect_num=0,
|
||
total_defect_area=0, rp=None)
|
||
print("图像中无果，跳过此图像")
|
||
continue
|
||
|
||
if cmd not in ['TO', 'PF', 'YR', 'KO']:
|
||
logging.error(f'错误指令，指令为{cmd}')
|
||
continue
|
||
|
||
spec = None
|
||
if cmd == 'PF':
|
||
start_time2 = time.time()
|
||
spec_data = pipe.receive_spec_data(spec_receive)
|
||
_, spec = pipe.parse_spec(spec_data)
|
||
end_time2 = time.time()
|
||
print(f'接收光谱数据时间：{end_time2 - start_time2}秒')
|
||
|
||
start_time3 = time.time()
|
||
if images: # 确保images不为空
|
||
response = process_data(cmd, images, spec, dp, pipe, detector)
|
||
end_time3 = time.time()
|
||
print(f'处理时间：{end_time3 - start_time3}秒')
|
||
if response:
|
||
logging.info(f'处理成功，响应为: {response}')
|
||
else:
|
||
logging.error('处理失败')
|
||
else:
|
||
print("没有有效的图像进行处理")
|
||
|
||
end_time = time.time()
|
||
print(f'全流程时间：{end_time - start_time}秒')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main(is_debug=False)
|