TensorFlow+Pytorch识别阿猫阿狗

猫狗大战

前言

这个是一次大作业，然后最近花了两三天把它训练完并且搭建起了可以用的服务。

作业内容就是猫狗大战（猫狗数据集分类），要求是用tensorflow和pytorch分别实现。这本来是几年前kaggle中的一个竞赛，原本数据集有800多M，但是我为了省训练时间，从网上找了一个“阉割版”的数据集，一共就3000张图片。具体的下载方式在后面的代码中会提到。

环境的话，限制于显卡算力，Pytorch版本还是比较低的

python3.8

TF==2.2

Pytorch==1.2

1. TensorFlow版

import tensorflow as tf
import os
import random
from tensorflow.keras import models,layers
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense,Conv2D,Flatten,Dropout,MaxPool2D
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.callbacks import EarlyStopping
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix

plt.rcParams['font.sans-serif'] = ['simhei']
plt.rcParams['axes.unicode_minus'] = False
复制代码

1.1 获取数据集

#第一次运行需要将注释取消进行数据下载

# 数据集下载链接
#dataset_url = "https://storage.googleapis.com/mledu-datasets/cats_and_dogs_filtered.zip"

# 开始下载，并且解压提取到设定文件夹
#dataset_path = tf.keras.utils.get_file("cats_and_dogs_filtered.zip", origin=dataset_url,cache_subdir="/home/a/桌面/python相关/人工智能期末作业-猫狗大战分类",extract=True)

dataset_dir = os.path.join(os.path.dirname('/home/a/桌面/python相关/人工智能期末作业-猫狗大战分类/'), "cats_and_dogs_filtered")
复制代码

第一次运行的时候将注释的几行代码取消注释就可以下载数据集了。

1.2 载入划分训练集，并且构造数据生成器

# 将下载的数据按标签划分路径
train_cats = os.path.join(dataset_dir,"train","cats")
train_dogs = os.path.join(dataset_dir,"train","dogs")

test_cats = os.path.join(dataset_dir,"validation","cats")
test_dogs = os.path.join(dataset_dir,"validation","dogs")

train_dir = os.path.join(dataset_dir,"train")
test_dir = os.path.join(dataset_dir,"validation")


# 查看数据大小
train_dogs_num = len(os.listdir(train_dogs))
train_cats_num = len(os.listdir(train_cats))

test_dogs_num = len(os.listdir(test_dogs))
test_cats_num = len(os.listdir(test_cats))

train_all = train_cats_num+train_dogs_num
test_all = test_cats_num+test_dogs_num


print(train_all,test_all)
复制代码

训练集有2000张，测试集1000张，这样的话训练集与测试集比例为2：1其实是不太好的，往往我们还是喜欢7/3开，不过也很接近了。

在构造数据生成器之前，或者说进行任何数据操作之前，我们必须要查看数据，对数据有一定了解。通过查看一些图片，发现每一张图片的大小并不一致，那么就要想一个比较好的大小来约束所有图片(在保证显存训练够用的情况下)。下面就是构造生成器的一些处理

• 设置batch_size
• 读取文件夹下的图片
• 为了降低计算量以及数据大小，先将图片RGB归一化，变为0~1之间
• 设置好图片的大小，打乱数据
• 设置好seed，保证可重现
• 设置分类指标（二分类）

batch_size=64
height=224
width=224

train_generator=ImageDataGenerator(
    rescale=1./255.
).flow_from_directory(
    batch_size=batch_size,
    directory=train_dir,
    shuffle=True,
    seed=0,
    target_size=(height,width),
    class_mode="binary"
)


test_generator=ImageDataGenerator(
    rescale=1./255.
).flow_from_directory(
    batch_size=batch_size,
    directory=test_dir,
    shuffle=False,
    seed=0,
    target_size=(height,width),
    class_mode="binary"
)
复制代码

然后利用构造好的生成器随机可视化一些图片

sample_training_images, labels = next(train_generator)
sample_testing_images,test_labels=next(test_generator)

d=train_generator.class_indices
names=dict(zip(d.values(),d.keys()))

def plotImages(images_arr,labels):
    fig, axes = plt.subplots(3, 5, figsize=(10,8))
    axes = axes.flatten()
    for (img,label), ax in zip(zip(images_arr,labels), axes):
        ax.imshow(img)
        ax.set_title("类别为: "+str(int(label))+" "+names[label])
        ax.axes.xaxis.set_visible(False)
        ax.axes.yaxis.set_visible(False)
    plt.tight_layout()
    plt.show()
plotImages(sample_training_images[:15],labels[:15])
plotImages(sample_testing_images[:15],test_labels[:15])
复制代码

1.3 模型构建与训练

一开始我满怀信心，打算构造一个复杂点的模型，直接搭建一个DenseNet

class ConvBlock(tf.keras.layers.Layer):
    def __init__(self, num_channels):
        super(ConvBlock, self).__init__()
        self.bn = tf.keras.layers.BatchNormalization()
        self.relu = tf.keras.layers.ReLU()
        self.conv = tf.keras.layers.Conv2D(
            filters=num_channels, kernel_size=(3, 3), padding='same')

        self.listLayers = [self.bn, self.relu, self.conv]

    def call(self, x):
        y = x
        for layer in self.listLayers.layers:
            y = layer(y)
        y = tf.keras.layers.concatenate([x,y], axis=-1)
        return y

    
    
# 输出通道数为num_convs*num_channels+输入通道数
class DenseBlock(tf.keras.layers.Layer):
    def __init__(self, num_convs, num_channels):
        super(DenseBlock, self).__init__()
        self.listLayers = []
        for _ in range(num_convs):
            self.listLayers.append(ConvBlock(num_channels))

    def call(self, x):
        for layer in self.listLayers.layers:
            x = layer(x)
        return x
    
    
class TransitionBlock(tf.keras.layers.Layer):
    def __init__(self, num_channels, **kwargs):
        super(TransitionBlock, self).__init__(**kwargs)
        self.batch_norm = tf.keras.layers.BatchNormalization()
        self.relu = tf.keras.layers.ReLU()
        self.conv = tf.keras.layers.Conv2D(num_channels, kernel_size=1)
        self.avg_pool = tf.keras.layers.AvgPool2D(pool_size=2, strides=2)

    def call(self, x):
        x = self.batch_norm(x)
        x = self.relu(x)
        x = self.conv(x)
        return self.avg_pool(x)
        
def block_1():
    return tf.keras.Sequential([
        tf.keras.layers.Conv2D(256, kernel_size=7, strides=2, padding='same'),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.ReLU(),
        tf.keras.layers.MaxPool2D(pool_size=3, strides=2, padding='same')])

def block_2():
    net = block_1()
    # num_channels为当前的通道数
    num_channels, growth_rate = 256,32
    num_convs_in_dense_blocks = [4,4,4,4]

    for i, num_convs in enumerate(num_convs_in_dense_blocks):
        net.add(DenseBlock(num_convs, growth_rate))
        # 上一个稠密块的输出通道数
        num_channels += num_convs * growth_rate
        # 在稠密块之间添加一个转换层，使通道数量减半
        if i != len(num_convs_in_dense_blocks) - 1:
            num_channels //= 2
            net.add(TransitionBlock(num_channels))
    return net


def DenseNet():
    net = block_2()
    net.add(tf.keras.layers.BatchNormalization())
    net.add(tf.keras.layers.GlobalAvgPool2D())
    net.add(tf.keras.layers.LeakyReLU(0.1))
    net.add(tf.keras.layers.Flatten())
    net.add(tf.keras.layers.Dense(128))
    net.add(tf.keras.layers.LeakyReLU(0.1))
    net.add(tf.keras.layers.Dense(1,activation='sigmoid'))
    return net
复制代码

训练试一下

num_epochs=5
lr=1e-4

# 实例化网络
densenet=DenseNet()

optimizer=tf.keras.optimizers.RMSprop(lr=lr)
loss=tf.keras.losses.BinaryCrossentropy(from_logits=True)

densenet.build(input_shape=(None,height,width,3))
densenet.summary()

densenet.compile(optimizer=optimizer,loss=loss,metrics=['accuracy'])
history=densenet.fit(
    train_generator,
    epochs=num_epochs,
    validation_data=test_generator,
)

densenet.evaluate(test_generator)
复制代码

def trainning_plot(history,num_epochs):
    x=[i for i in range(num_epochs)]
    plt.figure()
    plt.plot(x,history.history['accuracy'],label='accuracy')
    plt.plot(x,history.history['val_accuracy'],label='val_accuracy')
    plt.plot(x,history.history['loss'],label='loss')
    plt.plot(x,history.history['val_loss'],label='val_loss')
    plt.legend()
    plt.xlabel("Epochs")
    plt.show()
trainning_plot(history,num_epochs)
复制代码

从结果上看，模型实际上和盲猜的准确率一样，也就是说什么也没有学到。这也就是由于训练集和测试集比例不当造成的，总体训练的数据相对来说还是比较少。因此，为了扩充数据集，增强模型的鲁棒性和泛化能力，需要使用数据增强。

# 加入对训练数据进行平移、旋转、随机缩放等等
train_generator=ImageDataGenerator(
    rescale=1./255.,
    rotation_range=40,
    width_shift_range=0.1,
    height_shift_range=0.1,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True,
    fill_mode='nearest',
).flow_from_directory(
    batch_size=batch_size,
    directory=train_dir,
    shuffle=True,
    seed=0,
    target_size=(height,width),
    class_mode="binary"
)
复制代码

再来训练试试

num_epochs=25

history=densenet.fit(
    train_generator,
    epochs=num_epochs,
    validation_data=test_generator
)
复制代码

有了一定的提升，但是对于二分类问题这种准确率还是太低了，并且训练曲线动荡。思考一下准确率较低的原因，往往是因为模型太过复杂而数据量较小导致的训练无法收敛。在这种情况下，那可以试试使用较为简单的模型进行训练，试试在简单模型下是否能较为平稳地训练，搭建了一个最基本的CNN

from tensorflow.keras import layers,models,regularizers


model = models.Sequential()

model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(224, 224, 3)))
model.add(layers.MaxPooling2D((2, 2)))

model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))

model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))

model.add(layers.Conv2D(128, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))

model.add(layers.Conv2D(128, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))

model.add(layers.Flatten())
model.add(layers.Dense(512, activation='relu', kernel_regularizer=regularizers.l2(0.001)))
model.add(layers.Dropout(0.2))
model.add(layers.Dense(1, activation='sigmoid'))

print(model.summary())

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=5e-4), loss='binary_crossentropy', metrics='accuracy')
history = model.fit(
    train_generator,
    epochs=30,
    validation_data=test_generator,
    callbacks=[EarlyStopping(monitor='val_accuracy', min_delta=0.001, patience=5, verbose=1)]
)
复制代码

这里引入了正则化、早停、遗忘……希望能在简单模型上学习到比较好的参数

这里并没有早停而是训练了30个epoch，测试集上准确率提高到了81.6%，说明我们的猜想是对的（继续训练也许会继续提高准确率）。复杂的模型需要大量的数据来训练，而对于我们这少量的模型而言，简单的网络可能有更好的性能。既然如此，我们就可以利用迁移学习，将别人在大量数据集中训练好的模型拿过来直接用，这样也就避免了训练收敛困难的问题。

backbone=tf.keras.applications.DenseNet201(weights='imagenet',include_top=False,input_shape=(height,width,3))

backbone.trainable=False

transfer_model=Sequential()
transfer_model.add(backbone)
transfer_model.add(tf.keras.layers.GlobalAveragePooling2D())
transfer_model.add(Dense(512,activation='relu'))
transfer_model.add(Dense(1,activation='sigmoid'))

transfer_model.summary()

# 设置动态学习率，指数衰减
init_lr=1e-4
lr=tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate=init_lr,
    decay_steps=50,
    decay_rate=0.96,
    staircase=True
)

optimizer=tf.keras.optimizers.Adam(learning_rate=lr)
loss=tf.keras.losses.BinaryCrossentropy(from_logits=True)

transfer_model.compile(optimizer=optimizer, loss=loss, metrics='accuracy')
history = transfer_model.fit(
    train_generator,
    epochs=60,
    validation_data=test_generator,
    callbacks=[EarlyStopping(monitor='val_accuracy', min_delta=0.001, patience=5, verbose=1)]
)
复制代码

最终效果堪称完美，测试集上准确率达到98.9%，那就利用这个模型随机预测一些图片看看

plt.figure(figsize=(10,8))

# 获得原始的分类字典，并进行字典键值对互换
d=test_generator.class_indices
label_names=dict(zip(d.values(), d.keys()))


# 随机打乱测试集来看看预测效果
pre_generator=ImageDataGenerator(
    rescale=1./255.
).flow_from_directory(
    batch_size=batch_size,
    directory=test_dir,
    shuffle=True,
    seed=0,
    target_size=(height,width),
    class_mode="binary"
)

plt.suptitle("预测结果")
for images,labels in pre_generator:
    for i in range(25):
        ax = plt.subplot(5,5,i+1)
        plt.imshow(images[i])
        img_array = tf.expand_dims(images[i], 0)
        # 使用模型预测图片中的动物
        predictions = transfer_model.predict(img_array)
        predictions= 1 if predictions>=0.5 else 0
        plt.title(label_names[predictions])
        plt.axis("off")
    break
plt.show()
复制代码

看上去貌似都预测正确，再看看混淆矩阵

def plot_confusion_matrix(cm,classes, title='混淆矩阵'):

    plt.figure(figsize=(12, 8), dpi=100)
    np.set_printoptions(precision=2)

    # 在混淆矩阵中每格的概率值
    ind_array = np.arange(len(classes))
    x, y = np.meshgrid(ind_array, ind_array)
    for x_val, y_val in zip(x.flatten(), y.flatten()):
        c = cm[y_val][x_val]
        if c > 0.001:
            plt.text(x_val, y_val, "%0.2f" % (c,), color='red', fontsize=15, va='center', ha='center')
    
    plt.imshow(cm, interpolation='nearest')
    plt.title(title)
    xlocations = np.array(range(len(classes)))
    plt.xticks(xlocations, classes, rotation=90)
    plt.yticks(xlocations, classes)
    plt.ylabel('真实值')
    plt.xlabel('预测值')
    plt.show()
    
    
test_predict=transfer_model.predict_classes(test_generator,batch_size=batch_size)
test_names=list(test_generator.class_indices)

test_true=test_generator.classes

matrix=confusion_matrix(test_true,test_predict) 

plot_confusion_matrix(matrix,test_names)
复制代码

那这么好的模型当然得保存下来，留着后面用

transfer_model.save('tf_model/transfer_model')
复制代码

2. Pytorch版

大致思路和TensorFlow一样，只是两个框架不同api的使用上会有所区别。

import os
import torch
from torch import nn
from torch import functional as F
from torch.utils import data
from torchvision import transforms,datasets,models
import numpy as np
import time
import random
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix



plt.rcParams['font.sans-serif'] = ['simhei']
plt.rcParams['axes.unicode_minus'] = False
复制代码

2.1 载入数据

因为之前TensorFlow里已经下好了数据集，所以这里只用载入就好

file_path="./cats_and_dogs_filtered/"
train="train"
test="validation"


trans=transforms.Compose([
    transforms.Resize((224,224)),#随机切割将图片大小变为(224,224)
    transforms.ToTensor(), # 归一化为0-1
])


train_data=datasets.ImageFolder(os.path.join(file_path,train),trans)
test_data=datasets.ImageFolder(os.path.join(file_path,test),trans)


random_choice=random.sample([i for i in range(len(train_data))],25)
plt.figure(figsize=(10,8))
plt.suptitle("训练集可视化")
for i,j in enumerate(random_choice):
    ax = plt.subplot(5,5,i+1)
    plt.imshow(train_data[j][0].numpy().transpose((1,2,0)))
    plt.title("标签为: "+str(train_data[j][1])+" "+train_data.classes[train_data[j][1]])
    plt.axis("off")
plt.show()

batch_size=64

train_loader=data.DataLoader(train_data,batch_size=batch_size,shuffle=True)
test_loader=data.DataLoader(test_data,batch_size=batch_size,shuffle=False)
复制代码

在这里对图片处理的操作并不多，不像之前还有一系列图像增强操作。主要还是调整图片大小和归一化

2.2 模型构建与训练

有了之前的经验，我们干脆放弃花里胡哨的复杂网络，直接上基本的CNN

base_model=nn.Sequential(
    nn.Conv2d(3,48, kernel_size=7, stride=4, padding=1), nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Conv2d(48, 96, kernel_size=5, padding=2), nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Conv2d(96,128, kernel_size=3, padding=1), nn.ReLU(),
    nn.Conv2d(128,128, kernel_size=3, padding=1), nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Flatten(),
    nn.Linear(4608, 1024), nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(1024, 512), nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(512,2)
)


# 模型参数初始化
for name,param in base_model.named_parameters():
    if 'weight' in name:
        nn.init.kaiming_normal_(param)
    elif 'bias' in name:
        nn.init.constant_(param,val=0)
复制代码

各层输出如下

开始训练

epochs=40
lr=1e-4

cirterion=nn.CrossEntropyLoss()
optimizer=torch.optim.Adam(base_model.parameters(),lr=lr)

base_model=base_model.cuda()

base_model.train() 
loss_=0.
train_acc=0.
total=0.
for epoch in range(epochs):
    for i,data in enumerate(train_loader,0):
        inputs,train_labels=data
        optimizer.zero_grad()
        outputs=base_model(inputs.cuda())
        _,predicts=torch.max(outputs.data,1)
        train_acc+=(predicts.cuda()==train_labels.cuda().data).sum()
        loss=cirterion(outputs,train_labels.cuda())
        loss.backward()
        optimizer.step()
        
        loss_+=loss.item()
        #print(f"epoch: {epoch},loss: {loss_}")
        total+=train_labels.size(0)
    
    print(f"epoch: {epoch},loss={loss_/total*batch_size},acc={100*train_acc/total}%")  
复制代码

Pytorch训练起来比TensorFlow繁琐一些，需要在自己定义好损失函数以及优化器之后，进行正向传播计算损失、梯度，然后再利用优化器更新模型参数。当然，在高版本适配的Pytorch Lightning中已经实现了类似于TensorFlow.keras那样简单的compile、fit方法。不过我现在的显卡只能支持1.2，而Lightning最低要求1.3……为了使用GPU还是保持现状吧。

看看测试集上效果如何

def test(model,test_loader):
    model.eval()
    correct=0
    test_predict=[]
    with torch.no_grad():
        for idx,(t_data,t_target) in enumerate(test_loader):
            t_data,t_target=t_data.cuda(),t_target.cuda()
            pred=model(t_data)
            pred_class=pred.argmax(dim=1)
            test_predict.extend(pred_class.cpu())
            correct+=(pred_class==t_target).sum().item()
    acc=correct/len(test_data)
    print(f"测试集上准确率为: {acc*100}%")
    return test_predict

test_predict=test(base_model,test_loader)
复制代码

y_true=test_loader.dataset.targets

matrix=confusion_matrix(y_true,test_predict)


def plot_confusion_matrix(cm,classes, title='混淆矩阵'):

    plt.figure(figsize=(12, 8), dpi=100)
    np.set_printoptions(precision=2)

    # 在混淆矩阵中每格的概率值
    ind_array = np.arange(len(classes))
    x, y = np.meshgrid(ind_array, ind_array)
    for x_val, y_val in zip(x.flatten(), y.flatten()):
        c = cm[y_val][x_val]
        if c > 0.001:
            plt.text(x_val, y_val, "%0.2f" % (c,), color='red', fontsize=15, va='center', ha='center')
    
    plt.imshow(cm, interpolation='nearest')
    plt.title(title)
    xlocations = np.array(range(len(classes)))
    plt.xticks(xlocations, classes, rotation=90)
    plt.yticks(xlocations, classes)
    plt.ylabel('真实值')
    plt.xlabel('预测值')
    plt.show()
    
    
plot_confusion_matrix(matrix,list(test_loader.dataset.class_to_idx))
复制代码

效果和之前TensorFlow中的差不太多，那再试试迁移学习

transfer_model=models.densenet201(pretrained=True)
for param in transfer_model.parameters():
    param.requires_grad=False
    
transfer_model.classifier=nn.Sequential(
    nn.Linear(1920,512),
    nn.LeakyReLU(0.1),
    nn.Linear(512,128),
    nn.Dropout(0.5),
    nn.Linear(128,2)
)

transfer_model=transfer_model.cuda()


optimizer=torch.optim.Adam(transfer_model.parameters(),lr=lr)

epochs = 10

transfer_model.train()
loss_=0.
train_acc=0.
total=0.
for epoch in range(epochs):
    for i,data in enumerate(train_loader):
        inputs,train_labels=data
        optimizer.zero_grad()
        outputs=transfer_model(inputs.cuda())
        _,predicts=torch.max(outputs.data,1)
        train_acc+=torch.sum(predicts.cuda()==train_labels.cuda().data)
        loss=cirterion(outputs,train_labels.cuda())
        loss.backward()
        optimizer.step()
        
        loss_+=loss.item()
        #print(f"epoch: {epoch},loss: {loss_}")
        total+=train_labels.size(0)
    
    print(f"epoch: {epoch},loss={loss_/total*batch_size},acc={100*train_acc/total}%")  
复制代码

在迁移学习中，我们只需要将classifier层调整为我们自己需要的就好，前面的预训练模型不用动，也不用参加训练。

测试集上准确率依然有98.6%，也和之前TensorFlow上结果差不多，最后保存模型

torch.save(transfer_model,'./torch_model/transfer_model.pkl')
复制代码

3. 搭建图片分类服务

前面经过不断地调参、优化网络模型、尝试不同方法终于得到了性能优秀的模型，那就得好好利用起来。之前写过FastAPI的专栏，最后用sklearn做了一个demo。今天也和那个类似，使用本地保存下来的模型对上传的图片进行预测分类

# -*- coding: utf8 -*-
from PIL import Image
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.requests import Request
from fastapi.responses import RedirectResponse
from io import BytesIO
import tensorflow as tf
import uvicorn
import numpy as np
from typing import Optional, List
from starlette.templating import Jinja2Templates

tmp = Jinja2Templates(directory='templates')


class Model:
    model: Optional

    def load_model(self):
        self.model = tf.keras.models.load_model("./tf_model/transfer_model")

    def predict(self, input_image):
        output = self.model.predict_classes(input_image).item()
        mapping = {
            0: 'cat',
            1: 'dog'
        }

        return mapping[output]


def read_convert_image(file):
    loaded_image = Image.open(BytesIO(file))
    image_to_convert = np.asarray(loaded_image.resize((224, 224)))[..., :3]
    image_to_convert = np.expand_dims(image_to_convert, 0)
    image_to_convert = image_to_convert / 255.0
    return np.float32(image_to_convert)


describe = '''
<h2>访问/predict/image路由去尝试用训练好的模型对猫狗图片进行分类预测</h2>
'''
app = FastAPI(description=describe)
mymodel = Model()


@app.get("/predict/image")
def index(request: Request):
    return tmp.TemplateResponse('predict.html', {
        'request': request,
    })


@app.post("/predict/image")
async def image(request: Request, image_to_predict: UploadFile = File(...)):
    if image_to_predict is None or image_to_predict.file is None:
        raise HTTPException(status_code=400, detail="Please provide an image when calling this request")

    extension = image_to_predict.filename.split(".")[-1] in ("jpg", "jpeg", "png")
    if not extension:
        raise HTTPException(status_code=400, detail="Please provide an jpg or png image")
    img = image_to_predict.filename
    image_data = read_convert_image(image_to_predict.file.read())
    prediction = mymodel.predict(image_data)
    return tmp.TemplateResponse('result.html', {
        'request': request,
        "img": img,
        'prediction': prediction
    })


@app.get('/')
async def hello():
    return RedirectResponse("/docs")


@app.on_event("startup")
async def startup():
    mymodel.load_model()


if __name__ == "__main__":
    uvicorn.run("app:app", port=8000)

复制代码

依然是启动项目的时候加载模型，然后表单上传图片进行预测，返回结果。不过由于前端确实不好，想用html实现简单的页面渲染但是有的功能还是无法实现很好。无奈，只能做一个简陋的页面出来。

由于模型加载会花费较长时间，所以尽量别添加reload。下次试试TensorFlow serving

result.html

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>猫狗大战预测</title>
</head>
<body>
<h1>传入图片名称为:{{img}}</h1>

<h1>预测结果为:{{prediction}}</h1>

<a href="/predict/image"><strong>返回继续</strong></a>
</body>
</html>
复制代码

predict.html

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>猫狗大战预测</title>
</head>
<body>
<h1>上传一个猫/狗图片进行分类预测</h1>
<form action="/predict/image/" enctype="multipart/form-data" onchange="changepic(this)"  method="post">
    <input type="file" id="file" name="image_to_predict" accept="image/*">
    <input type="submit" value="预测">
</form>
<img src="" id="show" width="200">

</body>
<script>
    function changepic() {
        var reads= new FileReader();
        f=document.getElementById('file').files[0];
        reads.readAsDataURL(f);
        reads.onload=function (e) {
            document.getElementById('show').src=this.result;
        };
    }
</script>
</html>
复制代码

4. 最终效果

为了检验，去网上下载一些图试试