这是我完整的代码,杂乱的功能很多,实在抱歉,但为了代码完整性我全部传上来了。
代码是基于YOLO大作战来的,但是把pipeline的代码全部自己写(目的是获取sensor拍照获取的rgb888格式图片,但是比较笨只会这么写)。
第90行定义osd_img变量,用于储存yolo推理后显示推理框的图片。
第130行进行绘图yolo.draw_result(res,osd_img)
第133和134行进行图片格式转换和保存。
目前保存结果如下图,很不对劲。
希望大佬帮我看看怎样可以正常保存推理后显示推理框的图片。
from libs.YOLO import YOLO11
import os,sys,gc
import ulab.numpy as np
import nncase_runtime as nn
from media.sensor import *
from media.display import *
import image,time
from machine import FPIOA,Pin
def get_date_dir(str):
dir_name = f"/data/saveImg/{str}"
try:
os.mkdir(dir_name)
except OSError: # 文件夹已存在则跳过
pass
return dir_name
def destroy():
os.exitpoint(os.EXITPOINT_ENABLE_SLEEP)
# stop sensor
sensor.stop()
# deinit lcd
Display.deinit()
time.sleep_ms(50)
# deinit media buffer
MediaManager.deinit()
if __name__=="__main__":
fpioa = FPIOA()
fpioa.set_function(34,FPIOA.GPIO34)
fpioa.set_function(35,FPIOA.GPIO35)
valve_up=Pin(34,Pin.OUT,value = 0)
valve_low=Pin(35,Pin.OUT,value = 0)
up_count = 0
low_count = 0
counts = 4
isupdet = 0
islowdet = 0
#减少保存图片的次数
upsavecount=0
lowsavecount=0
#yolo推理存图
save_num = 0
# 显示模式,默认"hdmi",可以选择"hdmi"和"lcd"
display_mode="hdmi"
rgb888p_size=[640,360]
if display_mode=="hdmi":
display_size=[1920,1080]
#display_size=[640,360]
else:
display_size=[800,480]
kmodel_path = "/sdcard/best(1).kmodel"
labels = ["0"]
confidence_threshold = 0.2
nms_threshold=0.45
model_input_size=[320,320]
# 把pipeline的代码挪过来
# sensor给AI的图像分辨率
rgb888p_size=[ALIGN_UP(rgb888p_size[0],16),rgb888p_size[1]]
# 视频输出VO图像分辨率
display_size=[ALIGN_UP(display_size[0],16),display_size[1]]
# sensor对象
sensor = Sensor()
# osd显示Image对象
osd_img=None
nn.shrink_memory_pool()
# 初始化并配置sensor
sensor.reset()
# 通道0给到显示VO,格式为YUV420
sensor.set_framesize(w = display_size[0], h = display_size[1])
sensor.set_pixformat(PIXEL_FORMAT_YUV_SEMIPLANAR_420)
# 通道1用于保存图片,格式为RGB888,这个格式可以正常保存,PIXEL_FORMAT_RGB_888_PLANAR不行
sensor.set_framesize(w = rgb888p_size[0], h = rgb888p_size[1], chn=CAM_CHN_ID_1)
sensor.set_pixformat(Sensor.RGB888, chn=CAM_CHN_ID_1)
# 通道2给到AI做算法处理,格式为RGB888
sensor.set_framesize(w = rgb888p_size[0], h = rgb888p_size[1], chn=CAM_CHN_ID_2)
sensor.set_pixformat(PIXEL_FORMAT_RGB_888_PLANAR, chn=CAM_CHN_ID_2)
# OSD图像初始化
osd_img = image.Image(display_size[0], display_size[1], image.ARGB8888)
sensor_bind_info = sensor.bind_info(x = 0, y = 0, chn = CAM_CHN_ID_0)
Display.bind_layer(**sensor_bind_info, layer = Display.LAYER_VIDEO1)
# 初始化显示
if display_mode=="hdmi":
# 设置为LT9611显示,默认1920x1080
Display.init(Display.LT9611,osd_num=1, to_ide = True)
else:
# 设置为ST7701显示,默认480x800
Display.init(Display.ST7701, width=display_size[0], height=display_size[1], osd_num=1, to_ide=True)
# media初始化
MediaManager.init()
# 启动sensor
sensor.run()
save_dir34 = get_date_dir("IO34")
save_dir35 = get_date_dir("IO35")
# 初始化YOLOv11实例
yolo=YOLO11(task_type="detect",mode="video",kmodel_path=kmodel_path,labels=labels,rgb888p_size=rgb888p_size,model_input_size=model_input_size,display_size=display_size,conf_thresh=confidence_threshold,nms_thresh=nms_threshold,max_boxes_num=50,debug_mode=0)
yolo.config_preprocess()
try:
while True:
up_count = up_count - 1
low_count = low_count - 1
# 逐帧推理
save_src = sensor.snapshot(chn=CAM_CHN_ID_1)
ori = sensor.snapshot(chn=CAM_CHN_ID_2)
img=ori.to_numpy_ref()
# 获取图像的高度(假设 img 是一个 numpy 数组)
img_height = img.shape[1]
img_width = img.shape[2]
res=yolo.run(img)
yolo.draw_result(res,osd_img)
save_num +=1
img_test = osd_img.to_jpeg()
img_test.save(f"/sdcard/user_photos/{save_num}.jpg")
osd_img.draw_line(round(display_size[0]/2),0,round(display_size[0]/2),display_size[1],color = (255, 0, 0),thickness = 8)
osd_img.draw_string_advanced(0, 0, 48, "IO34")
osd_img.draw_string_advanced(round(display_size[0]/2), 0, 48, "IO35")
Display.show_image(osd_img, 0, 0, Display.LAYER_OSD3)
#pl.show_image()
# 遍历所有检测框
for detection in res:
# 提取检测框的坐标
x_min, y_min, x_max, y_max = detection[:4]
# 计算检测框的中心位置
center_x = (x_min + x_max) / 2
center_y = (y_min + y_max) / 2
# 判断中心位置是否在图像上半部分或下半部分
if center_x < img_width / 2:
up_count = up_count + 2
print("up_34") # 中心在上半部分
if isupdet == 0:
isupdet = 1
t1 = time.ticks_ms()
print("up_io34_first_detected")
else:
low_count = low_count + 2
print("low_35") # 中心在下半部
if islowdet == 0:
islowdet = 1
t2 = time.ticks_ms()
print("low_io35_first_detected")
print(f"up_count:", up_count)
print(f"low_count:", low_count)
if (up_count > counts):
valve_up.value(1)
up_count = 0
isupdet = 0
t3 = time.ticks_ms()
print(f"up_io34_found,time:{time.ticks_diff(t3, t1)}")
elif(up_count < -counts):
valve_up.value(0)
up_count = 0
isupdet = 0
print("up_io34_not_found")
timestamp = time.localtime()
time_part = "{:02d}{:02d}{:02d}".format(timestamp[3], timestamp[4], timestamp[5])
filename = f"{save_dir34}/{time_part}.rgb888"
upsavecount += 1
if upsavecount == 3:
save_src.save(filename)
upsavecount = 0
if (low_count > counts):
valve_low.value(1)
low_count = 0
islowdet = 0
t4 = time.ticks_ms()
print(f"low_io35_found:{time.ticks_diff(t4, t2)}")
elif(low_count < -counts):
valve_low.value(0)
low_count = 0
islowdet = 0
print("low_io35_not_found")
timestamp = time.localtime()
time_part = "{:02d}{:02d}{:02d}".format(timestamp[3], timestamp[4], timestamp[5])
filename = f"{save_dir35}/{time_part}.rgb888"
lowsavecount += 1
if lowsavecount == 3:
save_src.save(filename)
lowsavecount = 0
gc.collect()
except Exception as e:
#pass
print(f"Exception:{e}")
finally:
yolo.deinit()
destroy()