#!/usr/bin/env python3
"""라즈베리파이 추론 속도 실측.

사용:
python bench.py # 전체 조합 측정
python bench.py --img field.jpg # 실제 밭 사진으로 (권장)
python bench.py --soak 600 # 10분 지속 테스트 (발열 확인)
"""
import argparse, os, subprocess, time

torch와 ncnn이 각자 OpenMP 런타임을 올려 스레드를 중복 생성하면

힙이 깨진다 ("malloc(): unsorted double linked list corrupted" → Aborted).

2026-08-20 파이5에서 실제로 발생. 코어 수로 묶어 두면 사라진다.

★ torch/ultralytics import 전에 설정해야 한다. OpenMP는 로드 시점에 읽는다.

os.environ.setdefault("OMP_NUM_THREADS", str(os.cpu_count() or 4))

import numpy as np
from ultralytics import YOLO

def sh(cmd):
try:
return subprocess.check_output(cmd, shell=True).decode().strip()
except Exception:
return "?"

def temp():
return sh("vcgencmd measure_temp").replace("temp=", "")

def throttled():
"""0x0 이면 정상. 다른 값이면 전압 부족이나 발열로 성능이 깎인 것."""
return sh("vcgencmd get_throttled").replace("throttled=", "")

로봇 카메라가 실제로 내놓을 프레임 크기. 폰 원본(4032x3024)을 그대로 넣으면

매 추론마다 리사이즈가 섞여 15%쯤 느리게 나온다. 로봇은 그 일을 안 한다.

CAM_WH = (1280, 960)

def load_image(path, size=640):
if path:
import cv2
im = cv2.imread(path)
if im is None:
raise SystemExit(f"이미지를 못 읽었다: {path}")
if (im.shape[1], im.shape[0]) != CAM_WH:
print(f"입력 {im.shape[1]}x{im.shape[0]} -> {CAM_WH[0]}x{CAM_WH[1]} 로 맞춤")
im = cv2.resize(im, CAM_WH)
return im

# 사진이 없으면 노이즈로 대신한다 (속도는 거의 같다)
return np.random.randint(0, 255, (size, size, 3), dtype=np.uint8)

def bench(model, img, imgsz, budget=15.0, warmup=3):
"""budget 초 동안 계속 돌려서 평균 fps를 구한다."""
for _ in range(warmup):
model.predict(img, imgsz=imgsz, verbose=False)
n, t0 = 0, time.time()
while time.time() - t0 < budget:
model.predict(img, imgsz=imgsz, verbose=False)
n += 1
dt = time.time() - t0
return n / dt, dt / n * 1000.0

def run_matrix(img, budget):
combos = [
("yolo11n.pt", 416), ("yolo11n.pt", 640),
("yolo11s.pt", 416), ("yolo11s.pt", 640),
("yolo11n_ncnn_model", 416), ("yolo11n_ncnn_model", 640),
("yolo11s_ncnn_model", 416), ("yolo11s_ncnn_model", 640),
]
print(f"\n시작 온도 {temp()} throttled {throttled()}")
print(f"\n{'모델':<24}{'해상도':>8}{'fps':>9}{'ms/장':>10}{'온도':>10}")
print("-" * 61)
for name, imgsz in combos:
try:
m = YOLO(name)
except Exception as e:
print(f"{name:<24}{imgsz:>8} 건너뜀 ({type(e).name})")
continue
fps, ms = bench(m, img, imgsz, budget)
print(f"{name:<24}{imgsz:>8}{fps:>9.2f}{ms:>10.1f}{temp():>10}")
print(f"\n끝 온도 {temp()} throttled {throttled()}")
print("throttled 가 0x0 이 아니면 전원 부족이나 발열로 성능이 깎인 것 — 그 숫자는 못 쓴다")

def runsoak(model_name, imgsz, img, seconds):
"""오래 돌리면서 성능이 떨어지는지 본다. 여름 하우스 대비."""
m = YOLO(model_name)
print(f"\n지속 테스트 {model_name} @ {imgsz} {seconds}초")
print(f"\n{'경과(분)':>9}{'fps':>9}{'온도':>10}{'throttled':>12}")
print("-" * 40)
t_end = time.time() + seconds
while time.time() < t_end:
fps, = bench(m, img, imgsz, budget=30, warmup=0)
elapsed = seconds - (t_end - time.time())
print(f"{elapsed/60:>9.1f}{fps:>9.2f}{temp():>10}{throttled():>12}")

if name == "main":
p = argparse.ArgumentParser()
p.add_argument("--img", default=None, help="테스트용 사진 경로 (권장: 실제 밭 사진)")
p.add_argument("--budget", type=float, default=15.0, help="조합당 측정 시간(초)")
p.add_argument("--soak", type=int, default=0, help="지속 테스트 시간(초). 예: 600")
p.add_argument("--model", default="yolo11n_ncnn_model", help="지속 테스트용 모델")
p.add_argument("--imgsz", type=int, default=416, help="지속 테스트용 해상도")
a = p.parse_args()

image = load_image(a.img)
if a.soak:
    run_soak(a.model, a.imgsz, image, a.soak)
else:
    run_matrix(image, a.budget)