基于yolov5和desnet的貓咪識別模型

前言

前段時間給學校的貓咪小程序搭建了識貓模型，可以通過貓咪的照片辨別出是那只貓貓，這里分享下具體的方案，先看效果圖：

源代碼在文末

模型訓練

在訓練服務器（或你的個人PC）上拉取本倉庫代碼。

圖片數據準備

進入`data`目錄，執行`npm install`安裝依賴。（需要 Node.js 環境，不確定老版本 Node.js 兼容性，建議使用最新版本。）

復制`config.demo.ts`文件并改名為`config.ts`，填寫Laf云環境的`LAF_APPID`；

執行`npm start`，腳本將根據小程序數據庫記錄拉取小程序云存儲中的圖片。

如果不打算從laf拉取數據，也可以自己制作數據集，只要保證文件格式如下就可以

catface文件下面的data文件中的photos中有若干個文件夾，每個文件夾名稱為id，文件夾下為圖片。

環境搭建

返回倉庫根目錄，執行`python -m pip install -r requirements.txt`安裝依賴。（需要Python>=3.8。不建議使用特別新版本的 Python，可能有兼容性問題。）

如果是linux系統，可以直接執行`bash prepare_yolov5.sh`拉取YOLOv5目標檢測模型所需的代碼，然后下載并預處理模型數據。如果是windows系統可以自己手動從gihub上拉取yolov5的模型。

執行`python3 data_preprocess.py`，腳本將使用YOLOv5從`data/photos`的圖片中識別出貓貓并截取到`data/crop_photos`目錄。

開始訓練

執行`python3 main.py`，使用默認參數訓練一個識別貓貓圖片的模型。（你可以通過`python3 main.py --help`查看幫助來自定義一些訓練參數。）程序運行結束時，你應當看到目錄的export文件夾下存在`cat.onnx`和`cat.json`兩個文件。（訓練數據使用TensorBoard記錄在`lightning_logs`文件夾下。若要查看準確率等信息，請自行運行TensorBoard。）

執行`python3 main.py --data data/photos --size 224 --name fallback`，使用修改后的參數訓練一個在YOLOv5無法找到貓貓時使用的全圖識別模型。程序運行結束時，你應當看到目錄的export文件夾下存在`fallback.onnx`和`fallback.json`兩個文件。

這里介紹下模型類的代碼，我們定義了學習率，網絡指定為densenet21

import torch
import torch.nn as nn
from torchvision import models
import torch.optim as optim
from pytorch_lightning import LightningModule
import torchmetrics
from typing import Tupleclass CatFaceModule(LightningModule):def __init__(self, num_classes: int, lr: float):super(CatFaceModule, self).__init__()self.save_hyperparameters()self.net = models.densenet121(num_classes=num_classes)self.loss_func = nn.CrossEntropyLoss()def forward(self, x: torch.Tensor) -> torch.Tensor:return self.net(x)def training_step(self, batch: Tuple[torch.Tensor, torch.LongTensor], batch_idx: int) -> torch.Tensor:loss, acc = self.do_step(batch)self.log('train/loss', loss, on_step=True, on_epoch=True)self.log('train/acc', acc, on_step=True, on_epoch=True)return lossdef validation_step(self, batch, batch_idx: int):loss, acc = self.do_step(batch)self.log('val/loss', loss, on_step=False, on_epoch=True)self.log('val/acc', acc, on_step=False, on_epoch=True)def do_step(self, batch: Tuple[torch.Tensor, torch.LongTensor]) -> Tuple[torch.Tensor, torch.Tensor]:# shape: x (B, C, H, W), y (B), w (B)x, y = batch# shape: out (B, num_classes)out = self.net(x)loss = self.loss_func(out, y)with torch.no_grad():# 每個類別分別計算準確率，以平衡地綜合考慮每只貓的準確率accuracy_per_class = torchmetrics.functional.accuracy(out, y, task="multiclass", num_classes=self.hparams['num_classes'], average=None)# 去掉batch中沒有出現的類別，這些位置為nannan_mask = accuracy_per_class.isnan()accuracy_per_class = accuracy_per_class.masked_fill(nan_mask, 0)# 剩下的位置取均值acc = accuracy_per_class.sum() / (~nan_mask).sum()return loss, accdef configure_optimizers(self) -> optim.Optimizer:return optim.Adam(self.parameters(), lr=self.hparams['lr'])

在模型訓練完畢后可以運行我編寫的modelTest，在這個文件中替換圖片為自己的圖片，觀察輸出是否正常,正常輸出是這樣的：

在這個輸出中，通過yolo檢測了圖片中是否含有貓咪，通過densenet對圖片所屬于的類進行概率計算，概率和id按照概率從大到小排序返回。

接口實現

我們訓練了兩個densenet模型，一個是全圖像的輸入為228的模型a，一個是輸入圖像為128的模型b，當請求打到服務器時，應用程序會先通過yolo檢測是否有貓，有的話就截取貓咪圖像，使用模型b；否則不截取，使用模型a。

以下是代碼：

from typing import Any
from werkzeug.datastructures import FileStorageimport torch
from PIL import Image
import numpy as np
import onnxruntime
from flask import Flask, request
from dotenv import load_dotenv
import os
import json
import time
from base64 import b64encode
from hashlib import sha256load_dotenv("./env", override=True)HOST_NAME = os.environ['HOST_NAME']
PORT = int(os.environ['PORT'])SECRET_KEY = os.environ['SECRET_KEY']
TOLERANT_TIME_ERROR = int(os.environ['TOLERANT_TIME_ERROR']) # 可以容忍的時間戳誤差(s)IMG_SIZE = int(os.environ['IMG_SIZE'])
FALLBACK_IMG_SIZE = int(os.environ['FALLBACK_IMG_SIZE'])CAT_BOX_MAX_RET_NUM = int(os.environ['CAT_BOX_MAX_RET_NUM']) # 最多可以返回的貓貓框個數
RECOGNIZE_MAX_RET_NUM = int(os.environ['RECOGNIZE_MAX_RET_NUM']) # 最多可以返回的貓貓識別結果個數print("==> loading models...")
assert os.path.isdir("export"), "*** export directory not found! you should export the training checkpoint to ONNX model."crop_model = torch.hub.load('yolov5', 'custom', 'yolov5/yolov5m.onnx', source='local')with open("export/cat.json", "r") as fp:cat_ids = json.load(fp)
cat_model = onnxruntime.InferenceSession("export/cat.onnx", providers=["CPUExecutionProvider"])with open("export/cat.json", "r") as fp:fallback_ids = json.load(fp)
fallback_model = onnxruntime.InferenceSession("export/cat.onnx", providers=["CPUExecutionProvider"])print("==> models are loaded.")app = Flask(__name__)
# 限制post大小為10MB
app.config['MAX_CONTENT_LENGTH'] = 10 * 1024 * 1024def wrap_ok_return_value(data: Any) -> str:return json.dumps({'ok': True,'message': 'OK','data': data})def wrap_error_return_value(message: str) -> str:return json.dumps({'ok': False,'message': message,'data': None})def check_signature(photo: FileStorage, timestamp: int, signature: str) -> bool:if abs(timestamp - time.time()) > TOLERANT_TIME_ERROR:return FalsephotoBase64 = b64encode(photo.read()).decode()photo.seek(0) # 重置讀取位置，避免影響后續操作signatureData = (photoBase64 + str(timestamp) + SECRET_KEY).encode()return signature == sha256(signatureData).hexdigest()@app.route("/recognizeCatPhoto", methods=["POST"])
@app.route("/recognizeCatPhoto/", methods=["POST"])
def recognize_cat_photo():try:photo = request.files['photo']timestamp = int(request.form['timestamp'])signature = request.form['signature']if not check_signature(photo, timestamp=timestamp, signature=signature):return wrap_error_return_value("fail signature check.")src_img = Image.open(photo).convert("RGB")# 使用 YOLOv5 進行目標檢測，結果為[{xmin, ymin, xmax, ymax, confidence, class, name}]格式results = crop_model(src_img).pandas().xyxy[0].to_dict('records')# 過濾非cat目標cat_results = list(filter(lambda target: target['name'] == 'cat', results))if len(cat_results) >= 1:cat_idx = int(request.form['catIdx']) if 'catIdx' in request.form and int(request.form['catIdx']) < len(cat_results) else 0# 裁剪出(指定的)catcat_result = cat_results[cat_idx]crop_box = cat_result['xmin'], cat_result['ymin'], cat_result['xmax'], cat_result['ymax']# 裁剪后直接resize到正方形src_img = src_img.crop(crop_box).resize((IMG_SIZE, IMG_SIZE))# 輸入到cat模型img = np.array(src_img, dtype=np.float32).transpose((2, 0, 1)) / 255scores = cat_model.run([node.name for node in cat_model.get_outputs()], {cat_model.get_inputs()[0].name: img[np.newaxis, :]})[0][0].tolist()# 按概率排序cat_id_with_score = sorted([dict(catID=cat_ids[i], score=scores[i]) for i in range(len(cat_ids))], key=lambda item: item['score'], reverse=True)else:# 沒有檢測到cat# 整張圖片直接resize到正方形src_img = src_img.resize((FALLBACK_IMG_SIZE, FALLBACK_IMG_SIZE))img = np.array(src_img, dtype=np.float32).transpose((2, 0, 1)) / 255scores = fallback_model.run([node.name for node in fallback_model.get_outputs()], {fallback_model.get_inputs()[0].name: img[np.newaxis, :]})[0][0].tolist()# 按概率排序cat_id_with_score = sorted([dict(catID=fallback_ids[i], score=scores[i]) for i in range(len(fallback_ids))], key=lambda item: item['score'], reverse=True)return wrap_ok_return_value({'catBoxes': [{'xmin': item['xmin'],'ymin': item['ymin'],'xmax': item['xmax'],'ymax': item['ymax']} for item in cat_results][:CAT_BOX_MAX_RET_NUM],'recognizeResults': cat_id_with_score[:RECOGNIZE_MAX_RET_NUM]})except BaseException as err:return wrap_error_return_value(str(err))if __name__ == "__main__":app.run(host=HOST_NAME, port=PORT, debug=False)

我們可以在本地運行，如果想測試的小伙伴可以把接口中密鑰校驗的代碼刪除，然后直接發送post請求即可。