import cv2
import numpy as np

# ------------------ 手动定义分类标签 ------------------
labels = ["noready", "ready"]  # 0: 盖板不到位, 1: 盖板到位

# ------------------ 配置 ------------------
IMG_SIZE = (640, 640)


def preprocess_cls(bgr_image_corrected, target_size=IMG_SIZE):
    """
    对已经校正方向（翻转后）的 BGR 图像进行预处理：
      - 转 RGB
      - Resize
      - 添加 batch 维度
    注意：此函数不再负责翻转！翻转应在调用前完成。
    """
    rgb = cv2.cvtColor(bgr_image_corrected, cv2.COLOR_BGR2RGB)
    resized = cv2.resize(rgb, target_size)
    input_tensor = np.expand_dims(resized, axis=0).astype(np.uint8)
    return input_tensor


def get_top1_class_str_from_output(outputs):
    if outputs is None or len(outputs) == 0:
        return "Error: No output"
    logits = outputs[0].flatten()
    top1_idx = int(np.argmax(logits))
    if 0 <= top1_idx < len(labels):
        return labels[top1_idx]
    else:
        return "Unknown"


def myFunc(rknn_lite, IMG):
    """
    统一推理接口（分类任务）
    - 输入原始 BGR 图像（可能倒置）
    - 内部先上下翻转（校正方向）
    - 模型在翻转后的图像上推理
    - 返回带结果的翻转后图像（即“正”的图像）
    """
    # 1. 上下翻转原始图像（校正摄像头倒装）
    corrected_img = cv2.flip(IMG, 0)  # 现在这是“正”的图像

    # 2. 预处理（使用校正后的图像）
    input_data = preprocess_cls(corrected_img, target_size=IMG_SIZE)

    # 3. 推理
    outputs = rknn_lite.inference(inputs=[input_data])

    # 4. 解析结果
    class_name = get_top1_class_str_from_output(outputs)

    # 5. 在校正后的图像上绘制结果
    vis_img = corrected_img.copy()
    text = f"Class: {class_name}"
    cv2.putText(
        vis_img,
        text,
        (10, 30),
        cv2.FONT_HERSHEY_SIMPLEX,
        1.0,
        (0, 255, 0),
        2,
        cv2.LINE_AA
    )

    return vis_img  # 返回的是上下翻转后的 BGR 图像（即视觉上“正”的图）