使用Flask(或FastAPI、Django等)框架建立后端系统,编写API接口,接收前端传递过来的手写数字图片的像素数据,调用训练好的模型进行预测,并将预测结果返回给前端
时间: 2024-02-19 19:02:57 浏览: 100
知识领域: 后端开发 技术关键词: Django、Flask、后端框架、RESTful API
以下是使用Flask框架建立后端系统,编写API接口,接收前端传递过来的手写数字图片的像素数据,调用训练好的模型进行预测,并将预测结果返回给前端的代码实现:
```python
from flask import Flask, request, jsonify
import numpy as np
import torch
import torchvision.transforms as transforms
app = Flask(__name__)
# 加载训练好的模型
model = torch.load('model.pth')
# 定义图像预处理函数
def preprocess(image):
transform = transforms.Compose([
transforms.Resize(28),
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
return transform(image).unsqueeze(0)
# 定义API接口
@app.route('/predict', methods=['POST'])
def predict():
# 从前端接收手写数字图片数据
image_data = request.files['image'].read()
image = Image.open(io.BytesIO(image_data))
# 进行图像预处理
preprocessed_image = preprocess(image)
# 调用训练好的模型进行预测
with torch.no_grad():
output = model(preprocessed_image)
_, predicted = torch.max(output.data, 1)
# 将预测结果返回给前端
return jsonify({'result': predicted.item()})
if __name__ == '__main__':
app.run()
```
在上述代码中,我们首先加载训练好的模型,并定义了一个图像预处理函数,用于将前端传递过来的手写数字图片进行预处理。然后,我们定义了一个API接口,使用POST方法接收前端传递过来的手写数字图片数据,进行预处理并调用训练好的模型进行预测,最后将预测结果返回给前端。在主函数中,我们使用Flask框架的run()方法来启动后端服务器。
阅读全文