import gradio as gr
import numpy as np
from sklearn import datasets
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from PIL import Image
import os

# 加载数据集
def load_dataset():
    data_dir = os.path.join(os.path.dirname(__file__), 'data')
    dataset = datasets.load_files(data_dir, load_content=True, recursive=False)
    return dataset

# 预处理数据
def preprocess_data(dataset):
    X = np.array([file[0] for file in dataset.images])
    y = np.array([file[1] for file in dataset.images])
    # 将标签编码为数字
    le = LabelEncoder()
    y = le.fit_transform(y)
    return X, y

# 训练模型
def train_model(X, y):
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42)
    model = RandomForestClassifier(n_estimators=100)
    model.fit(X_train, y_train)
    return model

# 加载数据集
dataset = load_dataset()
X, y = preprocess_data(dataset)
model = train_model(X, y)

# 定义一个函数来处理上传的图像并进行预测
def predict_image(image):
    # 将上传的图像转换为 numpy 数组
    image = np.array(image)
    # 这里需要添加您的图像预处理代码，例如调整大小、归一化等
    # 例如，如果您的模型需要 224x224 的输入，您可以这样做：
    image = image.resize((64, 64))  # 假设模型需要 64x64 的输入
    image_array = np.array(image).reshape(1, -1)  # 展平图像数组
    
    # 使用模型进行预测
    prediction = model.predict(image_array)
    
    # 返回预测结果
    return "猫" if prediction[0] == 0 else "狗"

# 创建 Gradio 接口
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(label="上传图片"),
    outputs="text",
    title="猫狗分类器",
    description="上传一张图片，模型将预测它是猫还是狗。"
)

# 启动应用
iface.launch()