#!/usr/bin/env python3
"""检测单张图片 - 通过 Unix socket 发送到推理服务"""
import sys
import socket
import struct
import json
import base64
import cv2

SOCK_PATH = "/tmp/edge-infer.sock"

def detect_image(image_path):
    img = cv2.imread(image_path)
    if img is None:
        print("无法读取图片:", image_path)
        return
    
    # 编码为 JPEG
    _, buf = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, 90])
    jpeg_bytes = buf.tobytes()
    
    # 连接推理服务
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.connect(SOCK_PATH)
    
    # 构造请求
    msg = {
        "stream_id": 0,
        "device_id": "cli-test",
        "url": "file://" + image_path,
        "ts": 1234567890.0,
        "jpeg_b64": base64.b64encode(jpeg_bytes).decode("utf-8")
    }
    
    # 发送
    data = json.dumps(msg).encode("utf-8")
    sock.sendall(struct.pack(">I", len(data)) + data)
    
    # 接收结果
    hdr = sock.recv(4)
    if not hdr:
        print("未收到响应")
        sock.close()
        return
    
    length = struct.unpack(">I", hdr)[0]
    result = json.loads(sock.recv(length).decode("utf-8"))
    sock.close()
    
    # 输出结果
    dets = result.get("detections", [])
    print(f"\n图片: {image_path}")
    print(f"尺寸: {img.shape[1]}x{img.shape[0]}")
    print(f"检测到 {len(dets)} 个目标\n")
    print(f"{'类别':<20} {'置信度':<10} {'边界框'}")
    print("-" * 60)
    
    for d in dets:
        bbox = d["bbox"]
        print(f'{d["class"]:<20} {d["conf"]:<10.3f} [{bbox[0]:.0f}, {bbox[1]:.0f}, {bbox[2]:.0f}, {bbox[3]:.0f}]')
    
    # 保存带标注的图片
    if dets:
        for d in dets:
            x1, y1, x2, y2 = [int(x) for x in d["bbox"]]
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label = f'{d["class"]} {d["conf"]:.2f}'
            cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        out_path = image_path.rsplit(".", 1)[0] + "_result.jpg"
        cv2.imwrite(out_path, img)
        print(f"\n已保存标注图片: {out_path}")

if __name__ == "__main__":
    path = sys.argv[1] if len(sys.argv) > 1 else "/home/强光车灯误报.png"
    detect_image(path)