import gradio as gr
import cv2
import numpy as np
import pickle
import logging
import os


# 设置日志记录
def setup_logging():
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


# 加载模型函数
def load_model(model_filename):
    with open(model_filename, 'rb') as f:
        model = pickle.load(f)
    return model


# 图像预处理函数
def preprocess_image(image):
    # 检查图像是否有效
    if image is None or not isinstance(image, np.ndarray):
        raise ValueError("无效的图像输入")

    # 将图像转换为OpenCV格式
    img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    img = cv2.resize(img, (32, 32))  # 调整尺寸为32x32
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 转为灰度图像

    # 展平图像数据，并确保形状符合模型输入
    img_flattened = img.flatten().reshape(1, -1)  # 图像将展平为1024个特征
    return img_flattened


# 预测函数
def predict(image):
    try:
        img = preprocess_image(image)  # 预处理图像
        prediction = model.predict(img)  # 使用模型预测图像
        return "猫" if prediction[0] == 0 else "狗"
    except Exception as e:
        return f"处理图像时出错: {str(e)}"


setup_logging()


# 自动识别并加载最新的模型
def find_best_model():
    model_files = [f for f in os.listdir('.') if f.endswith('.pkl')]
    if not model_files:
        logging.error("未找到任何模型文件")
        return None
    best_model_filename = max(model_files, key=os.path.getmtime)
    logging.info(f"加载最新模型：{best_model_filename}")
    return best_model_filename


# 加载最佳模型
best_model_filename = find_best_model()
if best_model_filename:
    model = load_model(best_model_filename)

    # 创建 Gradio 界面
    iface = gr.Interface(
        fn=predict,
        inputs=gr.Image(type="numpy", label="上传图片"),
        outputs=gr.Textbox(label="预测结果"),
        live=True
    )
    iface.launch(share=True)  # 设置 share=True 以创建公共链接
else:
    logging.error("无法加载模型文件")