| from fastai.vision.all import * |
| import gradio as gr |
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| learn = load_learner("export_new.pkl") |
| |
| categories = tuple("数字" + str(x) for x in learn.dls.vocab) |
|
|
| def classify_image(img): |
| if isinstance(img, np.ndarray): |
| img = Image.fromarray(img) |
| if img.mode != 'L': |
| img = img.convert('L') |
| img = img.resize((28, 28)) |
| img = PILImage.create(img) |
| try: |
| pred, idx, probs = learn.predict(img) |
| print(f"预测类别: {pred}") |
| print(f"类别索引: {idx}") |
| print(f"概率分布: {probs}") |
| |
| if not torch.is_tensor(probs): |
| probs = torch.tensor(probs) |
| probs = F.softmax(probs, dim=0) |
| print("softmax后：", probs) |
| return dict(zip(categories, map(float, probs))) |
| except IndexError as e: |
| print(f"错误: {e}") |
| print(f"预测类别: {pred}", "检查模型的 vocab 是否匹配:", learn.dls.vocab) |
| error_msg = f"预测失败: {str(e)}" |
| raise RuntimeError(error_msg) from e |
|
|
| with gr.Blocks() as demo: |
| gr.Markdown("# Digits Classifier") |
| gr.Markdown("Classify digits numbers.") |
|
|
| with gr.Row(): |
| image_input = gr.Image(height=192, width=192) |
| label_output = gr.Label() |
|
|
| btn = gr.Button("Classify") |
| btn.click(fn=classify_image, inputs=image_input, outputs=label_output) |
|
|
| gr.Examples( |
| examples=[ |
| ["examples/number3.png"], |
| ["examples/number6.png"], |
| ["examples/number9.png"], |
| ], |
| inputs=image_input |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
