258 lines
8.7 KiB
Python
258 lines
8.7 KiB
Python
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
|
||
)
|
||
|
||
|