## 1.导入依赖包
```python
import io
import flask
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms as T
from torchvision.models import resnet50
```
## 2.初始化一个flask
```python
app = flask.Flask(__name__)
model = None
use_gpu = False # 是否使用GPU训练模型
with open('./data/class_map.txt', 'r') as f:
label_map = eval(f.read()) # 转化成字典
```
## 3.加载模型
```python
def load_model():
global model
model = resnet50(pretrained=True)
model.eval() # 不启用 BatchNormalization 和 Dropout
if use_gpu:
model.cuda() # 将模型加载到GPU上
```
## 4.处理接收到的图片
```python
def prepare_image(image, target_size):
if image.mode != 'RGB':
image = image.convert('RGB') # 使用'RGB'模式读取图片
image = T.Resize(target_size)(image)
image = T.ToTensor()(image)
image = T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(image)
image = image[None]
if use_gpu:
image = image.cuda()
return torch.autograd.Variable(image, volatile=True) # 自动微分变量
```
## 5.定义路由
```python
@app.route('/predict', methods=['POST'])
def predict():
data = {'success': False}
if flask.request.method == 'POST':
if flask.request.files.get('image'):
image = flask.request.files['image'].read()
image = Image.open(io.BytesIO(image)) # 将字节对象转为Byte字节流数据
image = prepare_image(image, target_size=(224, 224))
preds = F.softmax(model(image), dim=1)
results = torch.topk(preds.cpu().data, k=3, dim=1) # 返回Tensor中的前k个元素以及元素对应的索引值
results = (results[0].cpu().numpy(), results[1].cpu().numpy()) # 把tensor转换成numpy的格式
data['predictions'] = list()
for prob, label in zip(results[0][0], results[1][0]):
label_name = label_map[label]
r = {'label': label_name, 'probability': float(prob)}
data['predictions'].append(r)
data['success'] = True
return flask.jsonify(data) # 将字典转成json字符串
```
## 6.主函数
```python
if __name__ == '__main__':
load_model()
app.run()
```
