import torch from PIL import Image import torchvision.transforms as transforms import cv2 import time import os from model import resnet34 import serial from datetime import datetime import numpy as np import joblib import json import Find_COM from threading import Thread class MAT: def __init__(self, videoSourceIndex=0, weights_path = "resnet34-1Net.pth", json_path = 'class_indices.json', classes = 2): print('实验初始化中') self.data_root = os.getcwd() self.videoSourceIndex = videoSourceIndex # 摄像机编号 self.cap = cv2.VideoCapture(videoSourceIndex, cv2.CAP_DSHOW) # 打开摄像头 self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") self.port = Find_COM.list_ch340_ports()[0] # 串口名 self.pump_ser = serial.Serial(self.port, 9600) # 初始化串口 self.classes = classes self.total_volume = 0 # 记录总体积 self.now_volume = 0 # 记录当前注射泵内体积 self.volume_list = [] # 记录体积变化 self.voltage_list = [] # 记录电位变化（如有需要） self.color_list = [] # 记录颜色变化 self.start_time = time.time() # 记录实验开始时间 self.weights_path = os.path.join(self.data_root, weights_path) # 权重文件路径 self.json_path = os.path.join(self.data_root, json_path) # 类别文件路径 # 将开始时间转化为年月日时分秒的格式，后续文件命名都已此命名 self.formatted_time = datetime.fromtimestamp(self.start_time).strftime('%Y%m%d_%H%M%S') self.model = joblib.load("model.pkl") print("实验开始于", self.formatted_time) def get_picture(self, frame, typ=0, date=''): # 拍摄照片并保存 if frame is None: print(frame) image_name = f'{date}_{self.total_volume}.jpg' # 照片保存在Input文件夹下，以开始时间+体积数的方式命名 filepath = os.path.join(self.data_root, "Input", image_name) str_name = filepath.replace('%s', '1') cv2.imwrite(str_name, frame) return image_name def start_move_1(self): # 抽料程序 data = b"q1h24d" # *2 self.pump_ser.write(data) time.sleep(0.01) data = b"q2h0d" self.pump_ser.write(data) time.sleep(0.01) data = b"q4h0d" self.pump_ser.write(data) time.sleep(0.01) data = b"q5h15d" self.pump_ser.write(data) time.sleep(0.01) data = b"q6h3d" self.pump_ser.write(data) time.sleep(15) print('完成抽取') def start_move_2(self, speed=0.1): # 进料程序 # 计算单次滴定体积并传输至控制器 speed_min = speed * 30 speed_min_int = int(speed_min) speed_min_float = int((speed_min - speed_min_int) * 100) # print(speed_min_int, speed_min_float) data = f"q1h{speed_min_int}d" self.pump_ser.write(data.encode('ascii')) time.sleep(0.01) data = f"q2h{speed_min_float}d" self.pump_ser.write(data.encode('ascii')) time.sleep(0.01) data = b"q4h0d" self.pump_ser.write(data) time.sleep(0.01) data = b"q5h1d" self.pump_ser.write(data) time.sleep(0.01) # 进料 data = b"q6h2d" self.pump_ser.write(data) time.sleep(1) def start_move_3(self): # 进料急停 data = b"q6h6d" self.pump_ser.write(data) def preproc(self, im): try: hsv = cv2.cvtColor(im,cv2.COLOR_BGR2HSV) mask = hsv[:,:,1] > 150 mask = mask[:,:,np.newaxis] cnt = np.count_nonzero(mask) hsv*=mask h = round(np.sum(hsv[:,:,0])/cnt) s = round(np.sum(hsv[:,:,1])/cnt) v = round(np.sum(hsv[:,:,2])/cnt) return h,s,v except Exception as e : print(e) return None # name = f"{cl}_{h}_{s}_{v}.jpg" def _pred(self): suc,im = self.cap.read() if not suc: print("Failed to capture frame from camera.") return None ret = self.my_predictor(im) if ret is None: cv2.imwrite("tmp.jpg",im) return self.predictor("tmp.jpg") else: if ret == self.end_kind: print("Stop at ",self.total_volume) self.running = False self.start_move_3() else: self.thr = Thread(target=self._pred).start() return ret,0.9 def my_predictor(self,im): model = self.model ret = self.preproc(im) if ret is None: return None arr = np.array(ret) if len(arr.shape) == 1: arr = arr.reshape(1, -1) # 进行预测并转换为类别名 pred_labels = model.predict(arr) # pred_classes = [label_map[label] for label in pred_labels] mp = ["orange", "yellow"] if len(pred_labels) == 1: return mp[pred_labels[0]] else: return None def predictor(self, im_file): # 预测分类 image = Image.open(im_file) data_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) img = data_transform(image) img = torch.unsqueeze(img, dim=0) with open(self.json_path, "r") as f: class_indict = json.load(f) model = resnet34(num_classes=self.classes).to(self.device) assert os.path.exists(self.weights_path), "file: '{}' dose not exist.".format(self.weights_path) model.load_state_dict(torch.load(self.weights_path, map_location=self.device)) model.eval() with torch.no_grad(): output = torch.squeeze(model(img.to(self.device))).cpu() predict = torch.softmax(output, dim=0) predict_cla = torch.argmax(predict).numpy() class_a = "{}".format(class_indict[str(predict_cla)]) prob_a = "{:.3}".format(predict[predict_cla].numpy()) prob_b = float(prob_a) print('class_:',class_a) print('prob_:',prob_b) return class_a, prob_b def __del__(self): # 绘制滴定曲线 # self.line_chart() # 关闭串口和摄像头 self.pump_ser.close() self.cap.release() cv2.destroyAllWindows() print("Experiment finished.") def save_img(self): suc,im = self.cap.read() if not suc: print("Failed to capture frame from camera.") return cv2.imshow("new",im) name = f"Imgs/{self.formatted_time}_{self.total_volume}.jpg" if not cv2.imwrite(name,im): print("Failed to save image",name) def run(self,quick_speed = 0.2, mid_speed=0.1,slow_speed = 0.05,expect = 5, end_kind = 'orange', end_prob =0.5): n = 1 total_n = n # self.wait = False self.running = True self.end_kind = end_kind self.cnt = 0 self.thr = Thread(target=self._pred) self.thr.start() switching_point = expect * 0.9 while self.running: if self.now_volume <= 0: self.start_move_1() # 抽取12ml self.now_volume += 12 if self.total_volume < switching_point: # 每次加0.2ml speed = quick_speed self.start_move_2(speed) self.total_volume += speed self.now_volume -= speed else: speed = slow_speed self.start_move_2(speed) # 每次加0.05ml self.total_volume += speed self.now_volume -= speed self.total_volume = round(self.total_volume, 3) self.volume_list.append(self.total_volume) self.save_img() cv2.waitKey(1) print(f"Current Total Volume: {self.total_volume} ml") self.save_img() print('----->>Visual Endpoint<<-----') print(f"Total Volume: {self.total_volume} ml") # print(f"Image File: {im_file}") print("Volume List:", self.volume_list) print("Voltage List:", self.voltage_list) print("Color List:", self.color_list) if __name__ == "__main__": import warnings # 忽略所有警告 warnings.filterwarnings('ignore') # 创建MAT类的实例并运行 mat = MAT(videoSourceIndex = 0, weights_path = "resnet34-1Net.pth", json_path = 'class_indices.json', classes = 2) mat.run( quick_speed = 0.3, slow_speed = 0.2, expect = 11.2, end_kind = 'orange', end_prob = 0.5 )