模型地址:MaixHub

直装APP地址:MaixHub

APP对应的源码如下:


from maix import app, camera, comm, display, image, nn, time
import os
import struct


REPORT_ON = True
APP_CMD_DETECT_RES = 0x02

CAMERA_FPS = 60
CONFIDENCE_THRESHOLD = 0.60
NMS_IOU_THRESHOLD = 0.30
MIN_BOX_SIZE = 6
MIN_ASPECT_RATIO = 0.55
MAX_ASPECT_RATIO = 1.80
DUPLICATE_IOU_THRESHOLD = 0.35
CENTER_MERGE_RATIO = 0.50
COUNT_HISTORY_SIZE = 5
DRAW_INDEX = False


def encode_objs(objs):
    body = b""
    for obj in objs:
        body += struct.pack("<hhHHHf", obj.x, obj.y, obj.w, obj.h, obj.class_id, obj.score)
    return body


def box_iou(first, second):
    left = max(first.x, second.x)
    top = max(first.y, second.y)
    right = min(first.x + first.w, second.x + second.w)
    bottom = min(first.y + first.h, second.y + second.h)
    intersection = max(0, right - left) * max(0, bottom - top)
    if intersection <= 0:
        return 0.0

    union = first.w * first.h + second.w * second.h - intersection
    return intersection / union if union > 0 else 0.0


def is_duplicate(candidate, accepted):
    if box_iou(candidate, accepted) >= DUPLICATE_IOU_THRESHOLD:
        return True

    candidate_x = candidate.x + candidate.w * 0.5
    candidate_y = candidate.y + candidate.h * 0.5
    accepted_x = accepted.x + accepted.w * 0.5
    accepted_y = accepted.y + accepted.h * 0.5
    distance_squared = (candidate_x - accepted_x) ** 2 + (candidate_y - accepted_y) ** 2
    candidate_size = (candidate.w + candidate.h) * 0.5
    accepted_size = (accepted.w + accepted.h) * 0.5
    merge_distance = min(candidate_size, accepted_size) * CENTER_MERGE_RATIO
    return distance_squared < merge_distance * merge_distance


def filter_detections(objs):
    candidates = []
    for obj in objs:
        if obj.w < MIN_BOX_SIZE or obj.h < MIN_BOX_SIZE:
            continue
        aspect_ratio = obj.w / obj.h
        if MIN_ASPECT_RATIO <= aspect_ratio <= MAX_ASPECT_RATIO:
            candidates.append(obj)

    candidates.sort(key=lambda obj: obj.score, reverse=True)
    filtered = []
    for candidate in candidates:
        duplicate = False
        for accepted in filtered:
            if is_duplicate(candidate, accepted):
                duplicate = True
                break
        if not duplicate:
            filtered.append(candidate)
    return filtered


def update_stable_count(raw_count, history):
    history.append(raw_count)
    if len(history) > COUNT_HISTORY_SIZE:
        history.pop(0)
    ordered = sorted(history)
    return ordered[len(ordered) // 2]


model_path = "model_295047.mud"
if not os.path.exists(model_path):
    model_path = "/root/models/maixhub/295047/model_295047.mud"

detector = nn.YOLOv5(model=model_path)
cam = camera.Camera(
    detector.input_width(),
    detector.input_height(),
    detector.input_format(),
    fps=CAMERA_FPS,
)
dis = display.Display()
protocol = comm.CommProtocol(buff_size=1024)

count_history = []
frame_count = 0
fps_elapsed = 0
current_fps = 0.0

while not app.need_exit():
    frame_start = time.ticks_ms()
    img = cam.read()
    raw_objs = detector.detect(
        img,
        conf_th=CONFIDENCE_THRESHOLD,
        iou_th=NMS_IOU_THRESHOLD,
    )
    objs = filter_detections(raw_objs)
    raw_count = len(objs)
    ball_count = update_stable_count(raw_count, count_history)

    if REPORT_ON and objs:
        protocol.report(APP_CMD_DETECT_RES, encode_objs(objs))

    for index, obj in enumerate(objs, 1):
        img.draw_rect(obj.x, obj.y, obj.w, obj.h, color=image.COLOR_RED, thickness=2)
        if DRAW_INDEX:
            img.draw_string(
                obj.x,
                max(0, obj.y - 16),
                str(index),
                color=image.COLOR_RED,
                scale=0.8,
                thickness=1,
            )

    status = "BALLS:{}  RAW:{}  FPS:{:.1f}".format(ball_count, raw_count, current_fps)
    img.draw_string(9, 9, status, color=image.COLOR_BLACK, scale=1.25, thickness=3)
    img.draw_string(8, 8, status, color=image.COLOR_GREEN, scale=1.25, thickness=1)
    dis.show(img)

    frame_count += 1
    fps_elapsed += time.ticks_ms() - frame_start
    if frame_count % 30 == 0:
        current_fps = 30000.0 / fps_elapsed if fps_elapsed > 0 else 0.0
        fps_elapsed = 0

Logo

智能硬件社区聚焦AI智能硬件技术生态,汇聚嵌入式AI、物联网硬件开发者,打造交流分享平台,同步全国赛事资讯、开展 OPC 核心人才招募,助力技术落地与开发者成长。

更多推荐