在pycharm中利用h5模型调用摄像头识别多类水果并框选出水果位置以及标注水果名称
时间: 2023-11-18 14:02:26 浏览: 90
以下是通过Python和OpenCV实现的代码示例:
```python
import cv2
import numpy as np
import tensorflow as tf
# 加载模型
model = tf.keras.models.load_model('model.h5')
# 水果类别
classes = ['apple', 'banana', 'orange']
# 打开摄像头
cap = cv2.VideoCapture(0)
while True:
# 读取摄像头帧
ret, frame = cap.read()
# 调整大小
frame = cv2.resize(frame, (224, 224))
# 预处理图像
img = np.expand_dims(frame, axis=0)
img = img / 255.0
# 预测结果
pred = model.predict(img)
class_idx = np.argmax(pred)
class_name = classes[class_idx]
# 标注结果
cv2.putText(frame, class_name, (10, 30), cv2.FONT_HERSHEY_SIMPLEX,
1, (0, 255, 0), 2, cv2.LINE_AA)
# 框选水果位置
if class_idx == 0: # 苹果
lower = np.array([0, 50, 50])
upper = np.array([10, 255, 255])
elif class_idx == 1: # 香蕉
lower = np.array([20, 50, 50])
upper = np.array([30, 255, 255])
elif class_idx == 2: # 橙子
lower = np.array([10, 50, 50])
upper = np.array([20, 255, 255])
# 转换颜色空间
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
# 创建掩码
mask = cv2.inRange(hsv, lower, upper)
# 腐蚀和膨胀
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
mask = cv2.erode(mask, kernel, iterations=2)
mask = cv2.dilate(mask, kernel, iterations=2)
# 找到轮廓
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# 绘制边框
for c in contours:
x, y, w, h = cv2.boundingRect(c)
cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
# 显示结果
cv2.imshow('frame', frame)
# 按下q键退出
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# 释放摄像头并关闭窗口
cap.release()
cv2.destroyAllWindows()
```
在此示例中,我们首先加载了预训练的h5模型,然后打开了摄像头并读取每一帧。对于每一帧,我们将其调整到224x224大小并进行预处理。然后,我们将图像输入模型中进行预测,并找到最高概率的类别索引及其名称。
接下来,我们根据类别选择相应的颜色范围,并将图像转换到HSV颜色空间。然后,我们创建掩码并对其进行腐蚀和膨胀操作以去除噪声和填充水果区域。最后,我们找到水果区域的轮廓并绘制矩形边框。最后,我们将结果显示在窗口中。
请注意,此示例仅适用于较为简单的情况,例如水果在相机前方且没有重叠。如需更复杂的场景,请使用更复杂的算法。
阅读全文