251 lines
8.4 KiB
Python
251 lines
8.4 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
|
|
|
|
LOCAL_DEBUG = True
|
|
|
|
if LOCAL_DEBUG:
|
|
print("WARNING: Local debug mode is enabled. Serial communication will be skipped.")
|
|
time.sleep(2)
|
|
|
|
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")
|
|
if not LOCAL_DEBUG:
|
|
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.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 start_move_1(self): # 抽料程序
|
|
if LOCAL_DEBUG:return
|
|
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): # 进料程序
|
|
if LOCAL_DEBUG:
|
|
time.sleep(1)
|
|
return
|
|
# 计算单次滴定体积并传输至控制器
|
|
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): # 进料急停
|
|
if LOCAL_DEBUG:return
|
|
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.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'):
|
|
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',
|
|
)
|
|
|
|
|