addyya/animal_classification

1# 处理所有依赖模块的安装和导入2import os3import sys4from pathlib import Path5 6# 定义需要的模块列表7required_modules = [8    "fastai",9    "torch",10    "torchvision",11    "gradio",12    "pillow",13    "numpy"14]15 16# 安装缺失的模块17for module in required_modules:18    try:19        __import__(module)20    except ImportError:21        print(f"正在安装 {module}...")22        os.system(f"pip install {module}")23 24# 安装完成后导入所需模块25try:26    import torch27    from fastai.vision.all import *28    import gradio as gr29    from PIL import Image30    import numpy as np31except Exception as e:32    print(f"导入模块失败: {e}")33    sys.exit(1)34 35# 模型路径和错误处理36try:37    # 使用正斜杠路径（跨平台兼容）38    model_path = 'model/model.pkl'39    40    # 检查模型文件是否存在41    if not os.path.exists(model_path):42        print(f"模型文件不存在: {model_path}")43        possible_paths = [p for p in Path('.').rglob('*.pkl')]44        if possible_paths:45            print(f"发现可能的模型文件: {possible_paths}")46        sys.exit(1)47    48    # 加载模型49    learn = load_learner(model_path)50except Exception as e:51    print(f"加载模型失败: {e}")52    sys.exit(1)53 54# 分类函数55def is_cat(x): return x[0].isupper()56 57categories = ('Dog', 'Cat')58def classify_image(img):59    try:60        # 确保图像格式正确61        if not isinstance(img, PILImage):62            img = PILImage.create(img)63        64        # 预测65        pred, idx, probs = learn.predict(img)66        return dict(zip(categories, map(float, probs)))67    except Exception as e:68        print(f"分类错误: {e}")69        return {"error": str(e)}70 71# 测试图像72try:73    test_image_path = 'dog.jpg'74    if os.path.exists(test_image_path):75        im = PILImage.create(test_image_path)76        im.thumbnail((192, 192))77        print(f"测试分类结果: {classify_image(im)}")78    else:79        print(f"测试图像不存在: {test_image_path}")80except Exception as e:81    print(f"处理测试图像失败: {e}")82 83# 创建Gradio界面84try:85    with gr.Blocks() as demo:86        gr.Markdown("# 猫狗分类器")87        with gr.Row():88            with gr.Column():89                input_image = gr.Image(type="pil")90                submit_btn = gr.Button("分类")91            with gr.Column():92                output_label = gr.Label()93        94        submit_btn.click(fn=classify_image, inputs=input_image, outputs=output_label)95        96        gr.Examples(97            examples=["dog.jpg", "cat.jpg"] if os.path.exists("dog.jpg") else [],98            inputs=input_image99        )100    101    # 启动界面102    demo.launch(share=True)103except Exception as e:104    print(f"创建界面失败: {e}")105    sys.exit(1)