import os
import time
import ujson
import ustruct
import nncase_runtime as nn
import ulab.numpy as np
from machine import UART, FPIOA
from machine import Pin
from media.sensor import *
from media.display import *
from media.media import *
# 初始化FPIOA和GPIO
fpioa = FPIOA()
fpioa.set_function(53, FPIOA.GPIO53)
button = Pin(53, Pin.IN, Pin.PULL_DOWN) # 使用下拉电阻
# 显示模式配置
display_mode = "lcd" # "lcd" 或 "hdmi"
if display_mode == "lcd":
DISPLAY_WIDTH = ALIGN_UP(800, 16)
DISPLAY_HEIGHT = 480
else:
DISPLAY_WIDTH = ALIGN_UP(1920, 16)
DISPLAY_HEIGHT = 1080
# 图像尺寸配置
OUT_RGB888P_WIDTH = ALIGN_UP(1280, 16)
OUT_RGB888P_HEIGHT = 720
NUM_CLASSIFY_PATH = "/sdcard/num_classify/"
CONFIG_PATH = NUM_CLASSIFY_PATH + "deploy_config.json"
# 红色检测配置
RED_THRESHOLDS = [
(30, 100, 15, 127, 15, 127),
(30, 100, -127, -15, -127, -15)
]
BOXES = {
"B": (490, 0, 300, 720), # 蓝色方框
"R": (100, 210, 300, 300), # 红色方框
"Y": (880, 210, 300, 300) # 黄色方框
}
COLORS = {
"B": (0, 0, 255),
"R": (255, 0, 0),
"Y": (255, 255, 0)
}
# 初始化UART
fpioa.set_function(3, FPIOA.UART1_TXD)
fpioa.set_function(4, FPIOA.UART1_RXD)
uart = UART(UART.UART1, baudrate=115200)
def create_packet(x, y, z, n):
"""创建UART数据包"""
packet = bytearray(12)
mv = memoryview(packet) # 使用内存视图避免复制
# 帧头
mv[0] = 0x2C
mv[1] = 0x12
# 正确处理数字标签
try:
digit_value = int(n)
except (TypeError, ValueError):
digit_value = 0
# 数据部分:4个16位整数
ustruct.pack_into('>HHHH', mv, 2, x, y, z, digit_value)
# 计算校验和
checksum = 0
for i in range(2, 10):
checksum ^= mv[i]
mv[10] = checksum
# 帧尾
mv[11] = 0x5B
return packet
def read_deploy_config(config_path):
"""读取部署配置文件"""
with open(config_path, 'r') as json_file:
return ujson.load(json_file)
def softmax(x):
"""Softmax函数"""
exp_x = np.exp(x - np.max(x))
return exp_x / np.sum(exp_x)
def sigmoid(x):
"""Sigmoid函数"""
return 1 / (1 + np.exp(-x))
def main():
"""主函数"""
# 读取数字分类配置
deploy_conf = read_deploy_config(CONFIG_PATH)
kmodel_name = deploy_conf["kmodel_path"]
labels = deploy_conf["categories"]
confidence_threshold = deploy_conf["confidence_threshold"]
img_size = deploy_conf["img_size"]
num_classes = deploy_conf["num_classes"]
# 初始化KPU
kpu = nn.kpu()
kpu.load_kmodel(NUM_CLASSIFY_PATH + kmodel_name)
# 打印模型输入形状用于调试
try:
input_shape = kpu.get_input_shape(0)
print(f"模型输入形状: {input_shape}")
except Exception as e:
print(f"获取模型输入形状失败: {e}")
# 初始化AI2D预处理
ai2d = nn.ai2d()
ai2d.set_dtype(
nn.ai2d_format.NCHW_FMT,
nn.ai2d_format.NCHW_FMT,
np.uint8, np.uint8
)
ai2d.set_resize_param(True, nn.interp_method.tf_bilinear, nn.interp_mode.half_pixel)
# 确保AI2D输出形状与模型输入匹配
ai2d_output_shape = (1, 3, img_size[0], img_size[1])
ai2d_builder = ai2d.build([1, 3, OUT_RGB888P_HEIGHT, OUT_RGB888P_WIDTH],
ai2d_output_shape)
# 初始化传感器
sensor = Sensor()
sensor.reset()
sensor.set_hmirror(False)
sensor.set_vflip(False)
# 设置多路输出
sensor.set_framesize(width=DISPLAY_WIDTH, height=DISPLAY_HEIGHT, chn=CAM_CHN_ID_0) # 通道0: 显示
sensor.set_pixformat(PIXEL_FORMAT_YUV_SEMIPLANAR_420, chn=CAM_CHN_ID_0)
sensor.set_framesize(width=OUT_RGB888P_WIDTH, height=OUT_RGB888P_HEIGHT, chn=CAM_CHN_ID_1) # 通道1: 红色检测
sensor.set_pixformat(PIXEL_FORMAT_RGB_565, chn=CAM_CHN_ID_1)
sensor.set_framesize(width=OUT_RGB888P_WIDTH, height=OUT_RGB888P_HEIGHT, chn=CAM_CHN_ID_2) # 通道2: 数字分类
sensor.set_pixformat(PIXEL_FORMAT_RGB_888_PLANAR, chn=CAM_CHN_ID_2)
# 绑定显示
sensor_bind_info = sensor.bind_info(x=0, y=0, chn=CAM_CHN_ID_0)
if display_mode == "lcd":
Display.init(Display.ST7701, to_ide=True)
else:
Display.init(Display.LT9611, to_ide=True)
Display.bind_layer(**sensor_bind_info, layer=Display.LAYER_VIDEO1)
# 创建OSD层用于绘制结果
osd_img = image.Image(DISPLAY_WIDTH, DISPLAY_HEIGHT, image.ARGB8888)
# 初始化媒体
MediaManager.init()
sensor.run()
# 初始化AI2D输入输出张量
data = np.ones(ai2d_output_shape, dtype=np.uint8)
ai2d_output_tensor = nn.from_numpy(data)
# 添加数据包发送间隔控制变量
last_send_time = 0 # 上次发送时间戳
SEND_INTERVAL = 0.1 # 0.1秒发送间隔
# 主循环
clock = time.clock()
try:
while True:
clock.tick()
os.exitpoint()
# === 红色目标检测任务 ===
rgb565_img = sensor.snapshot(chn=CAM_CHN_ID_1)
if rgb565_img == -1:
print("获取通道1图像失败")
continue
red_in_blue = None # 确保在循环外部初始化
red_count = 0
red_blobs = [] # 初始化红色目标列表
for name, box in BOXES.items():
blobs = rgb565_img.find_blobs(
RED_THRESHOLDS,
False,
box,
x_stride=5,
y_stride=5,
pixels_threshold=1000,
area_threshold=1000,
merge=True,
margin=10
)
if blobs:
red_count += 1
if name == "B": # 蓝色方框
# 在蓝色方框中查找红色目标
red_blobs = rgb565_img.find_blobs(
RED_THRESHOLDS,
False,
box,
x_stride=5,
y_stride=5,
pixels_threshold=100,
area_threshold=100,
merge=True,
margin=10
)
# 处理红色目标检测结果
if red_blobs: # 检查列表是否非空
largest_blob = max(red_blobs, key=lambda b: b.pixels())
red_in_blue = (largest_blob.cx(), largest_blob.cy())
else:
red_in_blue = None
# === 数字分类任务 ===
# 核心修改:根据GPIO53状态决定是否进行数字识别
if button.value() == 1: # 高电平,进行数字识别
rgb888p_img = sensor.snapshot(chn=CAM_CHN_ID_2)
if rgb888p_img == -1:
print("获取通道2图像失败")
continue
cls_idx = -1
score = 0.0
digit_label = "0" # 默认值
if rgb888p_img.format() == image.RGBP888:
try:
ai2d_input = rgb888p_img.to_numpy_ref()
ai2d_input_tensor = nn.from_numpy(ai2d_input)
# 预处理并推理
ai2d_builder.run(ai2d_input_tensor, ai2d_output_tensor)
kpu.set_input_tensor(0, ai2d_output_tensor)
kpu.run()
# 获取输出
results = []
for i in range(kpu.outputs_size()):
output_data = kpu.get_output_tensor(i)
result = output_data.to_numpy()
results.append(result)
# 处理分类结果
if num_classes > 2:
softmax_res = softmax(results[0][0])
cls_idx = np.argmax(softmax_res)
if softmax_res[cls_idx] > confidence_threshold:
score = softmax_res[cls_idx]
digit_label = labels[cls_idx]
else:
sigmoid_res = sigmoid(results[0][0][0])
if sigmoid_res > confidence_threshold:
cls_idx = 1
score = sigmoid_res
digit_label = labels[1]
else:
cls_idx = 0
score = 1 - sigmoid_res
digit_label = labels[0]
except Exception as e:
print(f"数字分类处理失败: {e}")
else: # 低电平,不进行数字识别,直接设置为0
digit_label = "0"
cls_idx = -1
score = 0.0
# === UART数据包发送 ===
current_time = time.time()
# 检查是否达到发送间隔
if current_time - last_send_time >= SEND_INTERVAL:
try:
if red_in_blue:
x_val = red_in_blue[0]
y_val = red_in_blue[1]
else:
x_val, y_val = 0, 0
packet = create_packet(x_val, y_val, red_count, digit_label)
uart.write(packet) # 发送数据包
last_send_time = current_time # 更新最后发送时间
print(f"发送数据包: {packet}") # 打印发送日志
except Exception as e:
print(f"发送数据包失败: {e}")
# === 屏幕显示 ===
osd_img.clear()
# 绘制方框
for name, box in BOXES.items():
osd_img.draw_rectangle(box[0], box[1], box[2], box[3],
color=COLORS[name], thickness=2)
# 绘制红色目标坐标
if red_in_blue:
osd_img.draw_cross(red_in_blue[0], red_in_blue[1],
color=(255, 0, 0), size=8, thickness=2)
osd_img.draw_string_advanced(
BOXES["B"][0] + 5, BOXES["B"][1] + 5, 10,
f"Red: {red_in_blue[0]},{red_in_blue[1]}",
color=(255, 0, 0),
bg_color=(0, 0, 0)
)
# 显示红色目标数量
osd_img.draw_string_advanced(10, 10, 10, f"Red Count: {red_count}",
color=(255, 255, 0), bg_color=(0, 0, 0))
# 显示数字分类结果
button_state = "ON" if button.value() == 1 else "OFF"
osd_img.draw_string_advanced(
10, 40, 10,
f"Button: {button_state} | Digit: {digit_label}",
color=(0, 255, 0) if button.value() == 1 else (255, 0, 0),
bg_color=(0, 0, 0)
)
Display.show_image(osd_img, 0, 0, Display.LAYER_OSD3)
# === 控制台输出 ===
print(f"识别数字: {digit_label} (按钮状态: {button_state})")
print(f"红色目标数量: {red_count}")
if red_in_blue:
print(f"蓝色框内红色坐标: ({red_in_blue[0]}, {red_in_blue[1]})")
else:
print("蓝色框内红色坐标: 无")
except KeyboardInterrupt:
print("程序已停止")
except Exception as e:
print(f"主循环错误: {e}")
finally:
# 清理资源
try:
sensor.stop()
except Exception as e:
print(f"停止传感器失败: {e}")
try:
Display.deinit()
except Exception as e:
print(f"释放显示资源失败: {e}")
try:
MediaManager.deinit()
except Exception as e:
print(f"释放媒体资源失败: {e}")
try:
del ai2d_output_tensor
nn.shrink_memory_pool()
except Exception as e:
print(f"释放AI资源失败: {e}")
print("资源已释放")
if __name__ == "__main__":
main()
将显示通道0改为和红色检测通道1,其余不变
最新发布