From 1dea7ed92ba9f5d87043c52fa02ef3c60daf59eb Mon Sep 17 00:00:00 2001 From: FEIJINTI <83849113+FEIJINTI@users.noreply.github.com> Date: Fri, 16 Sep 2022 16:52:44 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E4=BA=86=E5=AF=B9=E5=BA=94?= =?UTF-8?q?=E7=9A=84=E9=80=9A=E8=AE=AF=E5=8D=8F=E8=AE=AE=EF=BC=88=E6=9C=AA?= =?UTF-8?q?=E6=B5=8B=E8=AF=95=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- socket_detector.py | 59 ++------------- utils.py | 183 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 188 insertions(+), 54 deletions(-) diff --git a/socket_detector.py b/socket_detector.py index a1487cb..f1c9e7a 100644 --- a/socket_detector.py +++ b/socket_detector.py @@ -4,65 +4,16 @@ import cv2 from classifer import WoodClass import time import os +from utils import PreSocket, process_cmd + def main(): - model_path = "models/model_2022-09-06_13-08.p" - socket_receive = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket_receive = PreSocket(socket.AF_INET, socket.SOCK_STREAM) socket_receive.connect(('127.0.0.1', 21122)) - socket_send = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket_send = PreSocket(socket.AF_INET, socket.SOCK_STREAM) socket_send.connect(('127.0.0.1', 21123)) - # _ = socket_receive.recv(4096*1000*10) - detector = WoodClass(w=4096, h=1200, n=3000, debug_mode=False) - detector.load(path=model_path) - while True: - # receive data - t1 = time.time() - size_buff = socket_receive.recv(5) - if (size_buff[4] == 0): - n_rows, n_cols = size_buff[0] << 8 | size_buff[1], size_buff[2] << 8 | size_buff[3] - data_size = n_rows * n_cols * 3 - print(data_size) - recv_size = data_size - buff_all, size = [], 0 - while True: - data_buff = socket_receive.recv(recv_size) - recv_size -= len(data_buff) - buff_all += data_buff - if recv_size == 0: - break - print(len(buff_all)) - raw_data = np.frombuffer(bytes(buff_all), dtype=np.uint8).reshape(int(n_rows), int(n_cols), -1) - print(raw_data.shape) - wood_color = detector.predict(raw_data) - # cv2.imshow("img", raw_data) - # cv2.waitKey(30) - # print('Class is ', wood_color) - if wood_color == 0: - socket_send.send(b'S') - print('S send success') - elif wood_color == 1: - socket_send.send(b'Z') - print('Z send success') - elif wood_color == 2: - socket_send.send(b'Q') - print('Q send success') - print((time.time()-t1)) - elif (size_buff[4] == 1): - detector = WoodClass(w=4096, h=1200, n=3000, debug_mode=False) - detector.correct() - detector.fit_pictures(data_path=r"C:\Users\FEIJINTI\PycharmProjects\wood_color") - socket_send.send(b'G') - elif (size_buff[4] == 2): - str_size = size_buff[0] << 8 | size_buff[1] - select_model = socket_receive.recv(str_size).decode("ascii") - model_path = os.path.join("models", select_model) - detector = WoodClass(w=4096, h=1200, n=3000, debug_mode=False) - detector.load(path=model_path) - socket_send.send(b'C') - - - + process_cmd(recv_sock=socket_receive , send_sock=socket_send) if __name__ == '__main__': diff --git a/utils.py b/utils.py index eec4a8f..6d14e85 100644 --- a/utils.py +++ b/utils.py @@ -5,9 +5,12 @@ Created on Nov 3 21:18:26 2020 @author: l.z.y @e-mail: li.zhenye@qq.com """ +import logging import os import shutil import time +import socket +from classifer import WoodClass def mkdir_if_not_exist(dir_name, is_delete=False): @@ -65,6 +68,186 @@ class Logger(object): print(content) +class PreSocket(socket.socket): + def __int__(self, *args, **kwargs): + super(socket.socket, self).__init__(*args, **kwargs) + self.pre_pack = b'' + + def receive(self, *args, **kwargs): + if len(self.pre_pack) == 0: + self.recv(*args, **kwargs) + else: + data_len = args[0] + required, left = self.pre_pack[:data_len], self.pre_pack[data_len:] + self.pre_pack = left + return required + + def set_prepack(self, pre_pack: bytes): + self.pre_pack += pre_pack + + +def process_cmd(recv_sock: PreSocket, send_sock: PreSocket): + model_path = "models/model_2022-09-06_13-08.p" + detector = WoodClass(w=4096, h=1200, n=3000, debug_mode=False) + detector.load(path=model_path) + while True: + pack, next_pack = receive_sock(recv_sock) + recv_sock.set_prepack(next_pack) + cmd, data = parse_protocol(pack) + ack_sock(send_sock, cmd_type=cmd) + if cmd == 'IM': + wood_color = str(detector.predict(data)) + done_sock(send_sock, cmd_type=cmd, result=wood_color) + elif cmd == 'TR': + detector.fit_pictures(data_path=r"C:\Users\FEIJINTI\PycharmProjects\wood_color") + done_sock(send_sock, cmd_type=cmd) + elif cmd == 'MD': + model_path = os.path.join("models", data) + detector.load(path=model_path) + done_sock(send_sock, cmd_type=cmd) + else: + logging.error(f'错误指令，指令为{cmd}') + + +def receive_sock(recv_sock: PreSocket, pre_pack: bytes = b'') -> (bytes, bytes): + """ + 从指定的socket中读取数据. + + :param recv_sock: 指定sock + :param pre_pack: 上一包的粘包内容 + :return: data, next_pack + """ + recv_sock.set_prepack(pre_pack) + # 开头校验 + while True: + try: + temp = recv_sock.receive(1) + except ConnectionError as e: + logging.error(f'连接出错, 错误代码:\n{e}') + return '', None + except TimeoutError as e: + logging.error(f'超时了，错误代码: \n{e}') + return '', None + except Exception as e: + logging.error(f'遇见未知错误，错误代码: \n{e}') + return '', None + if temp == b'aa': + break + + # 获取报文长度 + temp = b'' + while len(temp) < 4: + try: + temp += recv_sock.receive(1) + except Exception as e: + logging.error(f'接收报文长度失败, 错误代码: \n{e}') + return b'', b'' + try: + data_len = int.from_bytes(temp, byteorder='big') + except Exception as e: + logging.error(f'转换失败,错误代码 \n{e}, \n报文内容\n{temp}') + return b'', b'' + + # 读取报文内容 + temp = b'' + while len(temp) < data_len: + try: + temp += recv_sock.receive(data_len) + except Exception as e: + logging.error(f'接收报文内容失败, 错误代码: \n{e}，\n报文内容\n{temp}') + return b'', b'' + data, next_pack = temp[:data_len], temp[data_len:] + + # 进行数据校验 + temp = b'' + while len(temp) < 3: + try: + temp += recv_sock.receive(1) + except Exception as e: + logging.error(f'接收报文校验失败, 错误代码: \n{e}') + return b'', b'' + if temp == b'\xff\xff\xbb': + return data, next_pack + else: + logging.error(f"接收了一个完美的只错了校验位的报文，\n data: {data} \n next_pack:{next_pack}") + return b'', b'' + + +def parse_protocol(data: bytes) -> (str, any): + ''' + 指令转换. + + :param data:接收到的报文 + :return: 指令类型和内容 + ''' + try: + assert len(data) > 4 + except AssertionError: + logging.error('指令转换失败，长度不足5') + return '', None + cmd, data = data[:4], data[4:] + cmd = cmd.decode('ascii').strip().upper() + if cmd == 'IM': + n_rows, n_cols, img = data[:2], data[2:4], data[4:] + try: + n_rows, n_cols = [int.from_bytes(x, byteorder='big') for x in [n_rows, n_cols]] + except Exception as e: + logging.error(f'长宽转换失败, 错误代码{e}, 报文内容: n_rows:{n_rows}, n_cols: {n_cols}') + return '', None + try: + assert n_rows * n_cols * 3 == len(img) + except AssertionError: + logging.error('图像指令IM转换失败，数据长度错误') + return '', None + img = np.frombuffer(img, dtype=np.uint8).reshape((n_rows, n_cols, -1)) + return cmd, img + elif cmd == 'TR': + return cmd, None + elif cmd == 'MD': + data = data.decode('ascii') + return cmd, data + + +def ack_sock(send_sock:PreSocket, cmd_type: str) -> bool: + ''' + 发送应答 + :param cmd_type:指令类型 + :param send_sock:指定sock + :return:是否发送成功 + ''' + msg = b'\xaa\x00\x00\x00\xd8'+(' A'+cmd_type).upper().encode('ascii')+b'\xff\xff\xff\xbb' + try: + send_sock.send(msg) + except Exception as e: + logging.error(f'发送应答失败，错误类型：{e}') + return False + return True + + +def done_sock(send_sock: PreSocket, cmd_type: str, result: str = '') -> bool: + ''' + 发送任务完成指令 + :param cmd_type:指令类型 + :param send_sock:指定sock + :param result:数据 + :return:是否发送成功 + ''' + cmd_type = cmd_type.strip().upper() + if (cmd_type == "TR") or (cmd_type == "MD"): + if result != '': + logging.error('结果在这种指令里很没必要') + result = b'\xff' + elif cmd_type == 'IM': + result = result.encode('ascii') + msg = b'\xaa\x00\x00\x00\xd8'+(' D'+cmd_type).upper().encode('ascii') + result + b'\xff\xff\xbb' + try: + send_sock.send(msg) + except Exception as e: + logging.error(f'发送完成指令失败，错误类型：{e}') + return False + return True + + if __name__ == '__main__': log = Logger(is_to_file=True) log.log("nihao")