Can RK1820 Run YOLOv8? Let's Find Out

import cv2
import numpy as np
import time
from collections import deque
from rknn3lite.api.rknn3_lite import RKNN3Lite

MODEL = "yolov8_rk1820_fp16.rknn"
WEIGHT = "yolov8_rk1820_fp16.weight"

INPUT_SIZE = (640, 640)  
CONF_THRESH = 0.5
NMS_THRESH = 0.45
POSTPROCESS_INTERVAL = 2  

print("Loading RK1820 model...")
rknn = RKNN3Lite()
ret = rknn.load_rknn(
    model_path=MODEL,
    weight_path=WEIGHT
)

if ret != 0:
raise RuntimeError("Load model failed")

ret = rknn.init_runtime(
    target="rk1820",
    core_mask=0x01,
    device_id=None
)

if ret != 0:
raise RuntimeError("Init runtime failed")

print("RK1820 initialized.")

def preprocess_image(frame, target_size=INPUT_SIZE):

    img = cv2.resize(frame, target_size, interpolation=cv2.INTER_NEAREST)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    img = img.astype(np.float32) / 255.0

    img = np.expand_dims(img, axis=0)
    return img

def postprocess(outputs, frame_shape, conf_thresh=CONF_THRESH, nms_thresh=NMS_THRESH):

    if outputs is None or len(outputs) == 0:
        return [], 0.0
    
    output = outputs[0][0]  # [num_boxes, 6]
    
    scores = output[:, 4]
    valid_mask = scores >= conf_thresh
    output = output[valid_mask]
    scores = scores[valid_mask]
    
    if len(output) == 0:
        return [], 0.0
    
    boxes = output[:, :4].copy()
    class_ids = output[:, 5].astype(np.int32)
    
    h, w = frame_shape[:2]
    
    boxes[:, 0] = np.clip(boxes[:, 0], 0, w)
    boxes[:, 1] = np.clip(boxes[:, 1], 0, h)
    boxes[:, 2] = np.clip(boxes[:, 2], 0, w)
    boxes[:, 3] = np.clip(boxes[:, 3], 0, h)
    
    boxes_wh = np.zeros_like(boxes)
    boxes_wh[:, 0] = boxes[:, 0]
    boxes_wh[:, 1] = boxes[:, 1]
    boxes_wh[:, 2] = boxes[:, 2] - boxes[:, 0]
    boxes_wh[:, 3] = boxes[:, 3] - boxes[:, 1]
    
    # NMS
    indices = cv2.dnn.NMSBoxes(
        boxes_wh.tolist(),
        scores.tolist(),
        conf_thresh,
        nms_thresh
    )
    
    if len(indices) == 0:
        return [], 0.0
    
    if isinstance(indices, tuple):
        indices = indices[0]
    else:
        indices = indices.flatten()
    
    # 构建结果
    detections = []
    for idx in indices:
        detections.append({
            'box': boxes[idx].astype(np.int32),
            'score': float(scores[idx]),
            'class_id': int(class_ids[idx])
        })
    
    return detections, len(indices)

def draw_simple_boxes(frame, detections):
    """
    简化绘制：只画边框和标签，使用更快的绘制方式
    """
    colors = [
        (0, 255, 0), (255, 0, 0), (0, 0, 255),
        (255, 255, 0), (0, 255, 255), (255, 0, 255)
    ]
    
    for det in detections:
        box = det['box']
        x1, y1, x2, y2 = box
        score = det['score']
        class_id = det['class_id']
        
        color = colors[class_id % len(colors)]
        
        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 1)  # 线宽从 2 改为 1
        
        label = f"{score:.2f}"
        cv2.putText(
            frame, label,
            (x1, y1 - 5),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.4,  
            (0, 255, 0),
            1
        )
    
    return frame

cap = cv2.VideoCapture(0)

if not cap.isOpened():
raise RuntimeError("Cannot open camera")

cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)  

fps_buffer = deque(maxlen=30)
frame_count = 0
total_time = 0

cached_detections = []
frame_counter = 0

print("\n" + "=" * 60)
print(f"YOLOv8 Detection on RK1820 (Optimized)")
print(f"Input size: {INPUT_SIZE[0]}x{INPUT_SIZE[1]}")
print(f"Post-process every {POSTPROCESS_INTERVAL} frames")
print("Press 'q' or ESC to quit")
print("=" * 60 + "\n")

while True:
    t_start = time.perf_counter() 
    
    ret, frame = cap.read()
    if not ret:
        break
    
    frame_counter += 1
    img = preprocess_image(frame)
    
    outputs = rknn.inference(inputs=[img])
    
    if outputs is None:
        continue
    
    if frame_counter % POSTPROCESS_INTERVAL == 0 or frame_counter == 1:
        detections, num_boxes = postprocess(outputs, frame.shape)
        cached_detections = detections
    else:
        detections = cached_detections
        num_boxes = len(detections)
    
    if detections:
        frame = draw_simple_boxes(frame, detections)
    
    t_end = time.perf_counter()
    inference_time = t_end - t_start
    total_time += inference_time
    frame_count += 1
    
    fps = 1.0 / inference_time if inference_time > 0 else 0
    fps_buffer.append(fps)
    avg_fps = sum(fps_buffer) / len(fps_buffer)
    
    cv2.putText(
        frame,
        f"FPS: {avg_fps:.1f}",
        (10, 30),
        cv2.FONT_HERSHEY_SIMPLEX,
        1,
        (0, 255, 0),
        2
    )
    cv2.putText(
        frame,
        f"Boxes: {num_boxes}",
        (10, 60),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.6,
        (255, 255, 0),
        2
    )
    cv2.putText(
        frame,
        f"Res: {INPUT_SIZE[0]}x{INPUT_SIZE[1]}",
        (10, 85),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.5,
        (200, 200, 200),
        1
    )
    
    if frame_count % 10 == 0:
        print(f"Frame {frame_count}: FPS={avg_fps:.1f}, Boxes={num_boxes}")
    
    cv2.imshow("RK1820 Detection", frame)
    
    key = cv2.waitKey(1) & 0xFF
    if key == 27 or key == ord("q"):
        break
cap.release()
cv2.destroyAllWindows()
rknn.release()