问题描述
import time, os, gc, image
from media.sensor import *
from media.display import *
from media.media import *
from libs.YOLO import YOLO26
import _thread
DISPLAY_WIDTH = 800
DISPLAY_HEIGHT = 480
rgb888p_size = [320, 320]
tag_size = [160, 120]
display_size = [800, 480]
KMODEL_PATH = "/data/kmodel/best_yolo26_320.kmodel"
MODEL_INPUT_SIZE = [320, 320]
CONF_THRESH = 0.5
class_names = [
"0","1","2","3","4","5","6","7","8","9",
"10","11","12","13","14","15","16","17","18","19",
"20","21","22","23","24","25","26","27","28"
]
sensor = None
yolo_osd = None
tag_osd = None
lock = _thread.allocate_lock() # KPU inference lock only
yolo_stop = False
tag_stop = False
def yolo_thread():
global sensor, yolo_osd, yolo_stop
yolo = YOLO26(
task_type="detect", mode="video",
kmodel_path=KMODEL_PATH, labels=class_names,
rgb888p_size=rgb888p_size, model_input_size=MODEL_INPUT_SIZE,
display_size=display_size, conf_thresh=CONF_THRESH,
max_boxes_num=50, debug_mode=0
)
yolo.config_preprocess()
fps = 0.0
fc = 0
t0 = time.ticks_ms()
while not yolo_stop:
img = sensor.snapshot(chn=CAM_CHN_ID_2)
img_np = img.to_numpy_ref()
with lock:
res = yolo.run(img_np)
yolo_osd.clear()
if res is not None:
yolo.draw_result(res, yolo_osd)
fc += 1
now = time.ticks_ms()
dt = time.ticks_diff(now, t0)
if dt >= 1000:
fps = fc * 1000.0 / dt
fc = 0
t0 = now
yolo_osd.draw_string_advanced(8, 8, 24, "YOLO FPS:{:.1f}".format(fps), color=(255, 255, 0))
Display.show_image(yolo_osd, 0, 0, Display.LAYER_OSD1)
gc.collect()
yolo.deinit()
def apriltag_thread():
global sensor, tag_osd, tag_stop
last_id = -1
sx = display_size[0] / tag_size[0]
sy = display_size[1] / tag_size[1]
while not tag_stop:
img = sensor.snapshot(chn=CAM_CHN_ID_1)
tag_osd.clear()
tags = img.find_apriltags(families=image.TAG36H11)
for tag in tags:
last_id = tag.id()
r = tag.rect()
x = int(r[0] * sx); y = int(r[1] * sy)
w = int(r[2] * sx); h = int(r[3] * sy)
tag_osd.draw_rectangle(x, y, w, h, color=(255, 0, 0), thickness=2)
cx = int(tag.cx() * sx); cy = int(tag.cy() * sy)
tag_osd.draw_cross(cx, cy, color=(0, 255, 0), size=10, thickness=2)
tag_osd.draw_string_advanced(cx + 6, cy + 6, 24, "TAG:{}".format(last_id), color=(255, 255, 0))
tag_osd.draw_string_advanced(8, 34, 24,
"LAST TAG: {}".format(last_id) if last_id >= 0 else "LAST TAG: NONE",
color=(0, 255, 0) if last_id >= 0 else (200, 100, 100))
Display.show_image(tag_osd, 0, 0, Display.LAYER_OSD2)
gc.collect()
def media_init():
global sensor, yolo_osd, tag_osd
Display.init(Display.ST7701, width=DISPLAY_WIDTH, height=DISPLAY_HEIGHT, to_ide=True, osd_num=3)
sensor = Sensor(fps=30)
sensor.reset()
sensor.set_framesize(w=800, h=480, chn=CAM_CHN_ID_0)
sensor.set_pixformat(Sensor.YUV420SP)
sensor.set_framesize(w=tag_size[0], h=tag_size[1], chn=CAM_CHN_ID_1)
sensor.set_pixformat(Sensor.RGB565, chn=CAM_CHN_ID_1)
sensor.set_framesize(w=rgb888p_size[0], h=rgb888p_size[1], chn=CAM_CHN_ID_2)
sensor.set_pixformat(Sensor.RGBP888, chn=CAM_CHN_ID_2)
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)
yolo_osd = image.Image(display_size[0], display_size[1], image.ARGB8888)
tag_osd = image.Image(display_size[0], display_size[1], image.ARGB8888)
MediaManager.init()
sensor.run()
def media_deinit():
global sensor
os.exitpoint(os.EXITPOINT_ENABLE_SLEEP)
sensor.stop()
Display.deinit()
time.sleep_ms(50)
MediaManager.deinit()
if name == "main":
media_init()
_thread.start_new_thread(yolo_thread, ())
_thread.start_new_thread(apriltag_thread, ())
try:
while True:
time.sleep_ms(50)
except BaseException as e:
import sys
sys.print_exception(e)
yolo_stop = True
tag_stop = True
media_deinit()
gc.collect()
软件版本
CanMV_K230_01Studio_micropython_v1.6-b7e6201-dirty_nncase_v2.11.0.img.gz
