第134行的保存yolo处理后图像有问题,图像显示不对

Viewed 133

这是我完整的代码,杂乱的功能很多,实在抱歉,但为了代码完整性我全部传上来了。
代码是基于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()


2 Answers

使用下面的代码试试:

argb_np=pl.osd_img.to_numpy_ref()
rgb_np=argb_np[:,:,1:].copy()
img_565 = image.Image(rgb_np.shape[1], rgb_np.shape[0], image.RGB888, alloc=image.ALLOC_REF,data =rgb_np).to_rgb565()
img_565.save("/data/test.jpg")

如果想要保存带原画的结果,可以在snapshot的rgb888的图片上绘制,然后保存。这里需要区分RGB888P指的是数据排布为CHW,也就是RRRRRRR...GGGGGG...BBBBBB....,这个格式主要是为了给模型推理使用,因为模型输入通常是CHW的。而RGB888数据一般是HWC的,这是给Image类型使用的,排布是RGBRGBRGBRGB...,所以RGB888P无法直接用image.save保存。

argb_np=osd_img.to_numpy_ref()
rgb_np=argb_np[:,:,1:].copy()
img_565 = image.Image(rgb_np.shape[1], rgb_np.shape[0], image.RGB888, alloc=image.ALLOC_REF,data =rgb_np).to_rgb565()
img_565.save(f"/sdcard/user_photos/{save_num}.jpg")

我根据我的代码,省掉了"pl.",测试之后能保存图片,但是图片还是不对,图片效果全黑。我添加的字母和竖线却能正常显示

这是对的,osd_img本来就是个全透明的图,你绘制的只有结果,并不是在原图上绘制的。它和摄像头通道叠加才是你看到的效果