📚博客主頁:knighthood2001
?公眾號:認知up吧 (目前正在帶領大家一起提升認知,感興趣可以來圍觀一下)
🎃知識星球:【認知up吧|成長|副業】介紹
??如遇文章付費,可先看看我公眾號中是否發布免費文章??
🙏筆者水平有限,歡迎各位大佬指點,相互學習進步!
假設,你有模型,有訓練好的模型文件,有模型推理代碼,就可以把他放到flask上進行展示。
項目架構
index.html
是模板文件app.py
是項目運行的入口best_model.pth
是訓練好的模型參數model.py
是神經網絡模型,這里采用的是GoogleNet網絡。model_reasoning.py
是模型推理,通過這里面的代碼,我們可以在本地進行貓狗圖片的預測。
運行圖
點擊選擇文件
圖片下面就顯示預測結果了。
項目完整代碼與講解
index.html
<!DOCTYPE html>
<html lang="en">
<head><meta charset="UTF-8"><title>圖像分類</title><style>body {font-family: Arial, sans-serif;margin: 20px;}#result {margin-top: 10px;}#preview-image {max-width: 400px;margin-top: 20px;}</style>
</head>
<body><h1>圖像分類</h1><form id="upload-form" action="/predict" method="post" enctype="multipart/form-data"><input type="file" name="file" accept="image/*" onchange="previewImage(event)"><input type="submit" value="預測"></form><img id="preview-image" src="" alt=""><br><div id="result"></div><script>document.getElementById('upload-form').addEventListener('submit', async (e) => {e.preventDefault(); // 阻止默認的表單提交行為const formData = new FormData(); // 創建一個新的FormData對象,用于封裝表單數據formData.append('file', document.querySelector('input[type=file]').files[0]); // 添加表單數據// 使用fetch API發送POST請求到'/predict'路徑,并將formData作為請求體const response = await fetch('/predict', {method: 'POST',body: formData});// 獲取響應的JSON數據const result = await response.json();// 將預測結果顯示在頁面上ID為'result'的元素中document.getElementById('result').innerText = `預測結果: ${result.prediction}`;});function previewImage(event) {const file = event.target.files[0]; // 獲取上傳的文件對象const reader = new FileReader(); // 創建一個FileReader對象,用于讀取文件內容// 清空上一次的預測結果document.getElementById('result').innerText = '';// 當文件讀取完成后,將文件內容顯示在頁面上ID為'preview-image'的元素中reader.onload = function(event) {document.getElementById('preview-image').setAttribute('src', event.target.result);}// 如果用戶選擇了文件,則開始讀取文件內容if (file) {reader.readAsDataURL(file); // 將文件讀取為DataURL格式,這樣可以直接用作img元素的src屬性}}</script>
</body>
</html>
前端我練的不多,很多解釋已經在代碼中講了。
model.py
這是GoogleNet的網絡架構
import torch
from torch import nn
from torchsummary import summary
# 定義一個Inception模塊
class Inception(nn.Module):def __init__(self, in_channels, c1, c2, c3, c4): # 這些參數,所在的位置都會發送變化,所有需要這個參數super(Inception, self).__init__()self.ReLU = nn.ReLU()# 路線1,單1×1卷積層self.p1_1 = nn.Conv2d(in_channels=in_channels, out_channels=c1, kernel_size=1)# 路線2,1×1卷積層, 3×3的卷積self.p2_1 = nn.Conv2d(in_channels=in_channels, out_channels=c2[0], kernel_size=1)self.p2_2 = nn.Conv2d(in_channels=c2[0], out_channels=c2[1], kernel_size=3, padding=1)# 路線3,1×1卷積層, 5×5的卷積self.p3_1 = nn.Conv2d(in_channels=in_channels, out_channels=c3[0], kernel_size=1)self.p3_2 = nn.Conv2d(in_channels=c3[0], out_channels=c3[1], kernel_size=5, padding=2)# 路線4,3×3的最大池化, 1×1的卷積self.p4_1 = nn.MaxPool2d(kernel_size=3, padding=1, stride=1)self.p4_2 = nn.Conv2d(in_channels=in_channels, out_channels=c4, kernel_size=1)def forward(self, x):p1 = self.ReLU(self.p1_1(x))p2 = self.ReLU(self.p2_2(self.ReLU(self.p2_1(x))))p3 = self.ReLU(self.p3_2(self.ReLU(self.p3_1(x))))p4 = self.ReLU(self.p4_2(self.p4_1(x)))return torch.cat((p1, p2, p3, p4), dim=1)class GoogLeNet(nn.Module):def __init__(self, Inception, in_channels, out_channels):super(GoogLeNet, self).__init__()self.b1 = nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=7, stride=2, padding=3),nn.ReLU(),nn.MaxPool2d(kernel_size=3, stride=2, padding=1))self.b2 = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=64, kernel_size=1),nn.ReLU(),nn.Conv2d(in_channels=64, out_channels=192, kernel_size=3, padding=1),nn.ReLU(),nn.MaxPool2d(kernel_size=3, stride=2, padding=1))self.b3 = nn.Sequential(Inception(192, 64, (96, 128), (16, 32), 32),Inception(256, 128, (128, 192), (32, 96), 64),nn.MaxPool2d(kernel_size=3, stride=2, padding=1))self.b4 = nn.Sequential(Inception(480, 192, (96, 208), (16, 48), 64),Inception(512, 160, (112, 224), (24, 64), 64),Inception(512, 128, (128, 256), (24, 64), 64),Inception(512, 112, (128, 288), (32, 64), 64),Inception(528, 256, (160, 320), (32, 128), 128),nn.MaxPool2d(kernel_size=3, stride=2, padding=1))self.b5 = nn.Sequential(Inception(832, 256, (160, 320), (32, 128), 128),Inception(832, 384, (192, 384), (48, 128), 128),nn.AdaptiveAvgPool2d((1, 1)),nn.Flatten(),nn.Linear(1024, out_channels))for m in self.modules():if isinstance(m, nn.Conv2d):nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity='relu')if m.bias is not None:nn.init.constant_(m.bias, 0)elif isinstance(m, nn.Linear):nn.init.normal_(m.weight, 0, 0.01)if m.bias is not None:nn.init.constant_(m.bias, 0)def forward(self, x):x = self.b1(x)x = self.b2(x)x = self.b3(x)x = self.b4(x)x = self.b5(x)return xif __name__ == "__main__":device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = GoogLeNet(Inception, 1, 10).to(device)print(summary(model, (1, 224, 224)))
model_reasoning.py
import torch
from torchvision import transforms
from model import GoogLeNet, Inception
from PIL import Imagedef test_model(model, test_file):# 設定測試所用到的設備,有GPU用GPU沒有GPU用CPUdevice = "cuda" if torch.cuda.is_available() else 'cpu'model = model.to(device)classes = ['貓', '狗']print(classes)image = Image.open(test_file)# normalize = transforms.Normalize([0.162, 0.151, 0.138], [0.058, 0.052, 0.048])# # 定義數據集處理方法變量# test_transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor(), normalize])# 定義數據集處理方法變量test_transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])image = test_transform(image)# 添加批次維度,變成[1,3,224,224]image = image.unsqueeze(0)with torch.no_grad():model.eval()image = image.to(device) # 圖片也要放到設備當中output = model(image)print(output.tolist())pre_lab = torch.argmax(output, dim=1)result = pre_lab.item()print("預測值:", classes[result])return classes[result]def test_special_model(best_model_file, test_file):# 加載模型model = GoogLeNet(Inception, in_channels=3, out_channels=2)model.load_state_dict(torch.load(best_model_file))# 模型的推理判斷return test_model(model, test_file)if __name__ == "__main__":# # 加載模型# model = GoogLeNet(Inception, in_channels=3, out_channels=2)# model.load_state_dict(torch.load('best_model.pth'))# # 模型的推理判斷# test_model(model, "test_data/images.jfif")test_special_model("best_model.pth", "static/1.jpg")
這段代碼與之前的模型推理代碼不同的是,我添加了test_special_model
函數,方便后續app.py
中可以直接調用這個函數進行模型推理。
app.py
import os
from flask import Flask, request, jsonify, render_templatefrom model_reasoning import test_special_model
from model_reasoning import test_model
app = Flask(__name__)# 定義路由
@app.route('/')
def index():return render_template('index.html')@app.route('/predict', methods=['POST'])
def predict():if request.method == 'POST':# 獲取上傳的文件file = request.files['file']if file:# 調用模型進行預測# # 加載模型# model = GoogLeNet(Inception, in_channels=3, out_channels=2)# basedir = os.path.abspath(os.path.dirname(__file__))## model.load_state_dict(torch.load(basedir + '/best_model.pth'))# result = test_model(model, file)basedir = os.path.abspath(os.path.dirname(__file__))best_model_file = basedir + '/best_model.pth'result = test_special_model(best_model_file, file)return jsonify({'prediction': result})else:return jsonify({'error': 'No file found'})if __name__ == '__main__':app.run(debug=True)
如果沒有上文中的test_special_model
函數,那么這里你就需要
# 加載模型model = GoogLeNet(Inception, in_channels=3, out_channels=2)basedir = os.path.abspath(os.path.dirname(__file__))model.load_state_dict(torch.load(basedir + '/best_model.pth'))result = test_model(model, file)
并且還需要導入相應的庫。
best_model.pth
最重要的是,你需要訓練好的一個模型。
有需要的,可以聯系我,我直接把這個項目代碼發你。省得你還需要配置項目架構。
小插曲
我為什么會使用絕對路徑,因為我在使用相對路徑后,代碼提示找不到這個路徑。
basedir = os.path.abspath(os.path.dirname(__file__))best_model_file = basedir + '/best_model.pth'
然后,我剛剛又試了一下,發現使用相對路徑,又可以運行成功了。
真是不可思議(這個小插曲花了我大半個小時)。