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

with open('best_model.pkl', 'rb') as file:  # 替换为您的模型文件路径
    model = pickle.load(file)

def preprocess_image(image):
    img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # 将图片转换为灰度图像
    img = cv2.resize(img, (32, 32))  # 调整图片大小为 32x32 像素
    img = img.reshape(1, -1)  # 重塑为 1D 数组，如果模型需要 2D 数组，可以调整此行
    return img

def predict_image(image):
    processed_img = preprocess_image(image)
    prediction = model.predict(processed_img)
    return '猫' if prediction[0] == 0 else '狗'

iface = gr.Interface(
    fn=predict_image,
    inputs=gr.components.Image(), # 定义输入类型为图像
    outputs='text',
    title="猫狗图像分类器",
    description="上传一张图像来预测它是猫还是狗。"
)

iface.launch()