# 上料识别，获取上料图片，调用药材识别模型，输出识别结果
|
|
|
import os
|
import cv2
|
import time
|
import numpy as np
|
import onnxruntime
|
import win32com.client
|
|
|
|
# 计算拉普拉斯响应的方差,判断图像的清晰度
|
def variance_of_laplacian(image):
|
# 计算输入图像的拉普拉斯响应的方差
|
return cv2.Laplacian(image, cv2.CV_64F).var()
|
# 获取摄像头列表
|
def list_webcams():
|
# 创建一个WMI客户端实例
|
wmi = win32com.client.GetObject("winmgmts:")
|
webcams = wmi.InstancesOf("Win32_PnPEntity")
|
# 创建摄像头和索引字典
|
webcam_dict = {}
|
index = 0
|
for device in webcams:
|
name = getattr(device, 'Name', None)
|
pnp_class = getattr(device, 'PNPClass', None)
|
if name is not None and cam1 in name:
|
# 将设备名字和索引添加到字典中
|
webcam_dict[cam1] = index
|
index += 1
|
elif name is not None and cam2 in name:
|
# 将设备名字和索引添加到字典中
|
webcam_dict[cam2] = index
|
index += 1
|
return webcam_dict
|
|
|
# 调用另一个长焦镜头，拍摄清晰的局部药材图片
|
def get_image():
|
camera2_index = webcams.get(cam2)
|
print("第二个摄像头索引：" + str(camera2_index))
|
# 打开摄像头
|
capture = cv2.VideoCapture(camera2_index, cv2.CAP_DSHOW)
|
# 设置分辨率
|
capture.set(cv2.CAP_PROP_FRAME_WIDTH, 2048) # 宽度
|
capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 1540) # 高度
|
# 检查摄像头是否成功打开
|
if not capture.isOpened():
|
print("无法打开摄像头2")
|
exit()
|
width2 = capture.get(cv2.CAP_PROP_FRAME_WIDTH)
|
height2 = capture.get(cv2.CAP_PROP_FRAME_HEIGHT)
|
print("摄像头2分辨率:", width2, "x", height2)
|
# 循环读取摄像头画面
|
# Shadows name 'width' from outer scope
|
count = 0
|
while True:
|
ret2, frame2 = capture.read()
|
if not ret2:
|
print("无法读取摄像头画面")
|
break
|
|
count += 1
|
# 1920*1080的图像，中心裁剪640*480的区域
|
# a2 = int(height2 / 2 - target_height / 2)
|
# b2 = int(height2 / 2 + target_height / 2)
|
# c2 = int(width2 / 2 - target_width / 2)
|
# d2 = int(width2 / 2 + target_width / 2)
|
# cropped_frame2 = frame2[a2:b2, c2:d2]
|
if count == 2:
|
# 调整图像尺寸
|
resized_frame2 = cv2.resize(frame2, (target_width, target_height))
|
|
# 预处理
|
herb_blob = cv2.dnn.blobFromImage(resized_frame2, 1 / 255.0, (640, 640), swapRB=True, crop=False)
|
|
# 模型推理
|
herb_outputs = identify_session.run(None, {identify_session.get_inputs()[0].name: herb_blob})
|
|
herb_probabilities = herb_outputs[0]
|
|
top_five_classes = np.argsort(herb_probabilities, axis=1)[0][-5:][::-1]
|
name = ""
|
for i, class_id in enumerate(top_five_classes):
|
# 保留两位小数
|
probability = round(herb_probabilities[0][class_id], 2)
|
|
herb_class = herbs[class_id]
|
name = name + herb_class + "=" + str(probability)
|
# 显示画面
|
# cv2.imshow('Output2', resized_frame2)
|
|
# 计算拉普拉斯响应的方差
|
laplacian = variance_of_laplacian(resized_frame2)
|
# 生成保存文件名，以当前时间命名
|
save_name2 = name +"_["+ str(round(laplacian, 2)) +"]_"+ time.strftime("%Y%m%d%H%M%S", time.localtime()) + ".jpg"
|
# 判断图像的清晰度
|
# 保存调整尺寸后的图片
|
if laplacian > 1500:
|
c_ = save_path + "2/c/"
|
# 判断文件是否存在，不存在则创建
|
if not os.path.exists(c_):
|
os.makedirs(c_)
|
cv2.imwrite(c_ + save_name2, resized_frame2)
|
else:
|
n_ = save_path + "2/n/"
|
# 判断文件是否存在，不存在则创建
|
if not os.path.exists(n_):
|
os.makedirs(n_)
|
cv2.imwrite(n_ + save_name2, resized_frame2)
|
# cv2.imshow("Camera", resized_frame2)
|
print("保存图片:", save_name2)
|
break
|
# 结束线程
|
capture.release()
|
|
|
if __name__ == '__main__':
|
|
cam1 = "USB Camera"
|
cam2 = "PC Camera"
|
webcams = list_webcams()
|
print(webcams)
|
target_width = 1024
|
target_height = 768
|
save_path = "data/images/"
|
# 加载ONNX模型
|
session = onnxruntime.InferenceSession("model/loading.onnx")
|
identify_session = onnxruntime.InferenceSession("model/herb_identify.onnx")
|
herbs = eval(identify_session.get_modelmeta().custom_metadata_map['names'])
|
# 摄像头索引号，通常为0表示第一个摄像头
|
camera_index = webcams.get(cam1)
|
print("第一个摄像头索引：" + str(camera_index))
|
modelmeta = session.get_modelmeta()
|
metadata_map = modelmeta.custom_metadata_map
|
classes = eval(metadata_map['names'])
|
# 打开摄像头
|
cap = cv2.VideoCapture(camera_index, cv2.CAP_DSHOW)
|
# 设置分辨率
|
# cap.set(cv2.CAP_PROP_FRAME_WIDTH, 3840) # 宽度
|
# cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 2160) # 高度
|
# 检查摄像头是否成功打开
|
if not cap.isOpened():
|
print("无法打开摄像头")
|
exit()
|
|
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
|
height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
|
print("摄像头分辨率:", width, "x", height)
|
|
# 目标图像尺寸
|
|
# 计时器
|
frame_count = 0
|
start_time = time.time()
|
stime = time.time()
|
|
if not os.path.exists(save_path):
|
os.makedirs(save_path)
|
class_old = "1"
|
count = 0
|
status = "没有上料"
|
# 循环读取摄像头画面
|
while True:
|
# 睡眠100毫秒
|
time.sleep(0.1)
|
|
ret, frame = cap.read()
|
|
if not ret:
|
print("无法读取摄像头画面")
|
break
|
|
resized_frame = frame
|
if height > target_height and width > target_width:
|
# 1920*1080的图像，中心裁剪640*480的区域
|
a = int(height / 2 - target_height / 2)
|
b = int(height / 2 + target_height / 2)
|
c = int(width / 2 - target_width / 2)
|
d = int(width / 2 + target_width / 2)
|
print(a, b, c, d)
|
cropped_frame = frame[a:b, c:d]
|
# 调整图像尺寸
|
resized_frame = cv2.resize(cropped_frame, (target_width, target_height))
|
|
# 获取当前时间
|
current_time = time.time()
|
|
# 如果距离上一次保存已经过去1秒，则保存当前画面
|
# if current_time - start_time >= 3.0:
|
# # 生成保存文件名，以当前时间命名
|
# save_name = time.strftime("%Y%m%d%H%M%S", time.localtime()) + ".jpg"
|
# # 保存调整尺寸后的图片
|
# cv2.imwrite(save_path + save_name, frame)
|
# print("保存图片:", save_name)
|
# # 重置计时器
|
# start_time = time.time()
|
|
# 预处理
|
blob = cv2.dnn.blobFromImage(resized_frame, 1 / 255.0, (640, 640), swapRB=True, crop=False)
|
|
# 模型推理
|
outputs = session.run(None, {session.get_inputs()[0].name: blob})
|
#
|
# print(outputs)
|
# 应用softmax函数
|
probabilities = outputs[0]
|
|
# 找到最大概率的类别
|
predicted_class = np.argmax(probabilities, axis=1)[0]
|
max_probability = np.max(probabilities, axis=1)[0]
|
class_ = classes[predicted_class]
|
|
# 计算类型重复的次数，类别更换之后重新计数
|
if class_ != class_old:
|
count = 0
|
else:
|
count += 1
|
class_old = class_
|
print(f"{class_}:{count}: {max_probability}")
|
# 判断是否上料并且上料次数大于10次
|
if class_ == "shangliao" and count > 10:
|
status = "正在上料"
|
# 每隔3秒取一帧图像
|
# 如果距离上一次保存已经过去1秒，则保存当前画面
|
if current_time - stime >= 10.0:
|
save_name = time.strftime("%Y%m%d%H%M%S", time.localtime()) + ".jpg"
|
# 保存调整尺寸后的图片
|
path_ = save_path + "1/"
|
if not os.path.exists(path_):
|
os.makedirs(path_)
|
cv2.imwrite(path_ + save_name, resized_frame)
|
# 重置计时器
|
stime = time.time()
|
get_image()
|
else:
|
status = "没有上料"
|
|
print(status)
|
# # 找到概率较高的前十个类别
|
# top_ten_classes = np.argsort(probabilities, axis=1)[0][-1:]
|
#
|
# # 输出前十个类别
|
# print("Top 5 Classes:")
|
# for i in top_ten_classes:
|
# print(f"{classes[i]}: {probabilities[0][i]}")
|
|
# 计算帧速率
|
frame_count += 1
|
end_time = time.time()
|
elapsed_time = end_time - start_time
|
fps = frame_count / elapsed_time
|
# print(f"FPS: {fps:.2f}")
|
# 将FPS绘制在图像上
|
cv2.putText(resized_frame, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2,
|
cv2.LINE_AA)
|
# 显示画面
|
cv2.imshow("Camera", resized_frame)
|
|
# 检测按键，如果按下q键则退出循环
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
break
|
|
# 关闭摄像头
|
cap.release()
|
|
# 关闭所有窗口
|
cv2.destroyAllWindows()
|
