基于yolov5和desnet的猫咪识别模型

本文主要是介绍基于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请求即可。