# 从零到一:用YOLOv12打造你的专属宠物品种识别桌面应用
家里养了只猫,朋友来玩总爱问“这是什么品种?”,每次都要翻手机查半天。后来发现,很多宠物主人都有类似的困扰——面对市场上几十上百个猫狗品种,普通人很难一眼分辨清楚。要是能有个小工具,随手一拍就能告诉你宠物的品种,还能记录下识别历史,那该多方便。
这就是我们今天要动手实现的项目:一个完全由你掌控的宠物品种识别桌面应用。它不仅能识别37种常见的猫狗品种,还能通过摄像头实时检测、保存识别记录,甚至有个挺酷的科幻风格界面。最重要的是,整个过程从模型训练到界面开发,你都能亲手完成,不需要依赖任何云端服务,所有数据都留在本地。
如果你对Python有些基础,想尝试把深度学习模型变成真正可用的软件,或者单纯想给自家宠物做个有趣的小工具,这篇文章会带你走完全程。我会避开那些复杂的理论推导,聚焦在“怎么做”上,每个步骤都有可运行的代码,遇到坑的地方也会提前提醒。
## 1. 环境搭建与数据准备
开始之前,我们需要把开发环境搭建好。深度学习项目最怕的就是环境冲突,不同项目需要的库版本可能完全不同。我的习惯是每个项目都创建独立的虚拟环境,这样即使搞砸了也不会影响其他工作。
### 1.1 创建专属的Python环境
打开终端(Windows用命令提示符或PowerShell,macOS/Linux用终端),我们先用conda创建一个新的环境。如果你没有安装Anaconda,可以去官网下载Miniconda,它更轻量,完全够用。
```bash
# 创建名为pet_identifier的Python 3.9环境
conda create -n pet_identifier python=3.9 -y
# 激活环境
conda activate pet_identifier
```
激活后,你会看到命令行前面出现了`(pet_identifier)`,说明已经在这个环境里了。接下来安装PyTorch,这是我们的深度学习框架。这里有个小技巧:如果你有NVIDIA显卡并且想用GPU加速,需要安装CUDA版本的PyTorch;如果只是CPU运行,安装CPU版本就行。
```bash
# CPU版本(大多数人都能用)
pip install torch torchvision torchaudio
# 如果有NVIDIA显卡,去PyTorch官网根据你的CUDA版本选择安装命令
# 比如CUDA 11.8:
# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
```
> 注意:GPU版本能大幅提升训练和推理速度,但安装前要确认显卡驱动和CUDA版本。运行`nvidia-smi`可以查看CUDA版本,如果没输出说明可能没装驱动或者不是NVIDIA卡。
### 1.2 获取和处理宠物数据集
我们需要一个标注好的宠物数据集。网上能找到不少公开数据集,但质量参差不齐。我整理了一个包含37个品种、近1.4万张图片的数据集,涵盖了常见的12种猫和25种狗。数据已经按YOLO格式标注好了,你只需要下载解压就能用。
数据集结构是这样的:
```
pet_breed_dataset/
├── train/
│ ├── images/ # 训练图片
│ └── labels/ # 对应的标注文件
├── valid/ # 验证集
└── test/ # 测试集
```
每个标注文件是.txt格式,每行代表一个标注框,格式为:`类别ID x_center y_center width height`。所有坐标都是相对图片宽高的比例值(0到1之间)。
为了让YOLO能正确读取数据,我们需要创建一个配置文件`data.yaml`:
```yaml
# data.yaml
path: /path/to/your/pet_breed_dataset # 数据集根目录
train: train/images
val: valid/images
test: test/images
nc: 37 # 类别数量
names: [
'cat-Abyssinian', 'cat-Bengal', 'cat-Birman', 'cat-Bombay',
'cat-British_Shorthair', 'cat-Egyptian_Mau', 'cat-Maine_Coon',
'cat-Persian', 'cat-Ragdoll', 'cat-Russian_Blue', 'cat-Siamese',
'cat-Sphynx', 'dog-american_bulldog', 'dog-american_pit_bull_terrier',
'dog-basset_hound', 'dog-beagle', 'dog-boxer', 'dog-chihuahua',
'dog-english_cocker_spaniel', 'dog-english_setter',
'dog-german_shorthaired', 'dog-great_pyrenees', 'dog-havanese',
'dog-japanese_chin', 'dog-keeshond', 'dog-leonberger',
'dog-miniature_pinscher', 'dog-newfoundland', 'dog-pomeranian',
'dog-pug', 'dog-saint_bernard', 'dog-samoyed', 'dog-scottish_terrier',
'dog-shiba_inu', 'dog-staffordshire_bull_terrier',
'dog-wheaten_terrier', 'dog-yorkshire_terrier'
]
```
这个文件告诉YOLO数据在哪、有多少类别、每个类别叫什么名字。记得把`path`改成你实际存放数据集的路径。
## 2. YOLOv12模型训练与优化
有了数据,接下来就是训练模型。YOLOv12是Ultralytics公司发布的最新版本,在速度和精度上都有提升。更重要的是,它提供了从nano到large多个尺寸的预训练模型,我们可以根据需求选择。
### 2.1 选择合适的模型尺寸
YOLOv12有5个官方版本,特点对比如下:
| 模型尺寸 | 参数量 | 推理速度 | 适用场景 | 推荐设备 |
|---------|--------|----------|----------|----------|
| YOLOv12n | ~3M | 最快 | 移动端、嵌入式 | 树莓派、手机 |
| YOLOv12s | ~9M | 很快 | 实时检测 | 普通笔记本 |
| YOLOv12m | ~25M | 中等 | 平衡型应用 | 带GPU的PC |
| YOLOv12b | ~40M | 较慢 | 高精度需求 | 服务器 |
| YOLOv12l | ~60M | 最慢 | 研究、竞赛 | 多GPU工作站 |
对于宠物识别这种桌面应用,我推荐用YOLOv12s。它在普通笔记本上就能实时运行(30FPS以上),精度也足够识别大多数品种。如果你要在树莓派上跑,那就选YOLOv12n;如果追求极致精度而且有显卡,可以考虑YOLOv12m。
安装YOLO相关的库:
```bash
pip install ultralytics
```
这个命令会安装YOLOv12所需的所有依赖,包括OpenCV、Pillow等。
### 2.2 开始训练模型
训练代码比你想的简单得多。Ultralytics把复杂的训练过程封装成了几行代码:
```python
# train.py
from ultralytics import YOLO
def main():
# 加载预训练模型
model = YOLO('yolov12s.pt') # 会自动下载预训练权重
# 开始训练
results = model.train(
data='data.yaml', # 数据集配置文件
epochs=100, # 训练轮数
batch=8, # 批次大小
imgsz=640, # 输入图片尺寸
device='cpu', # 用CPU训练,如果是GPU改成'0'或'cuda'
workers=0, # 数据加载线程数(Windows建议设为0)
project='pet_train', # 保存结果的目录
name='exp1', # 实验名称
patience=20, # 早停耐心值
lr0=0.01, # 初始学习率
lrf=0.01, # 最终学习率因子
momentum=0.937, # 动量
weight_decay=0.0005, # 权重衰减
warmup_epochs=3, # 热身轮数
warmup_momentum=0.8, # 热身动量
box=7.5, # 框损失权重
cls=0.5, # 分类损失权重
dfl=1.5, # DFL损失权重
)
print("训练完成!")
print(f"最佳模型保存在: {results.save_dir}")
if __name__ == '__main__':
main()
```
运行这个脚本,训练就开始了。你会看到控制台输出类似这样的信息:
```
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
1/100 2.1G 1.234 2.567 1.891 32 640: 100%|██████████| 120/120 [01:23<00:00, 1.44it/s]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 15/15 [00:04<00:00, 3.21it/s]
all 368 1234 0.456 0.389 0.412 0.256
```
训练过程中有几个关键指标需要关注:
- **box_loss/cls_loss/dfl_loss**:损失值,应该随着训练逐渐下降
- **mAP50**:IoU阈值为0.5时的平均精度,我们的目标是在验证集上达到0.9以上
- **mAP50-95**:IoU阈值从0.5到0.95的平均精度,更严格的指标
> 提示:如果发现损失不下降或者精度很低,可能是学习率太大或太小。可以尝试把`lr0`调到0.001或0.0001。另外,数据量大的话可以增加`epochs`到150或200。
训练完成后,最佳模型会保存在`pet_train/exp1/weights/best.pt`。我们可以用这个模型进行测试:
```python
# test.py
from ultralytics import YOLO
import cv2
model = YOLO('pet_train/exp1/weights/best.pt')
# 测试单张图片
results = model('test_image.jpg', save=True, conf=0.5)
# 显示结果
for r in results:
im_array = r.plot() # 绘制检测框的图片
cv2.imshow('Result', im_array)
cv2.waitKey(0)
cv2.destroyAllWindows()
```
如果测试效果满意,就可以进入下一步——把模型封装成应用了。
## 3. 构建桌面应用界面
模型训练好了,但总不能每次都写代码来调用。我们需要一个图形界面,让不懂编程的人也能用。这里选择PyQt5,它是Python最成熟的GUI库之一,功能强大,文档丰富。
### 3.1 安装PyQt5和设计工具
```bash
pip install PyQt5 PyQt5-tools
```
PyQt5-tools里包含Qt Designer,这是一个可视化界面设计工具。虽然我们可以完全手写界面代码,但用Designer拖拽组件会快得多。
先设计主界面。打开Qt Designer(安装后在开始菜单找,或者运行`designer.exe`),创建一个Main Window。我设计了一个科幻风格的界面,主要包含这些区域:
1. **顶部工具栏**:模型选择、参数调节
2. **中央显示区**:左右并排显示原图和检测结果
3. **底部结果表格**:显示检测到的品种、置信度、位置
4. **侧边控制面板**:开始/停止检测、保存结果等按钮
设计完后保存为`main_window.ui`,然后用pyuic5转换成Python代码:
```bash
pyuic5 main_window.ui -o ui_main.py
```
生成的`ui_main.py`包含了界面布局的所有代码,但我们不能直接修改它(因为每次重新生成都会覆盖)。正确的做法是创建一个新类来继承它:
```python
# main_window.py
import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QMessageBox, QFileDialog
from PyQt5.QtCore import QThread, pyqtSignal, Qt
from PyQt5.QtGui import QImage, QPixmap
import cv2
import numpy as np
from ultralytics import YOLO
from ui_main import Ui_MainWindow # 自动生成的界面类
class DetectionThread(QThread):
"""专门处理检测任务的线程,避免界面卡顿"""
frame_processed = pyqtSignal(np.ndarray, np.ndarray, list)
finished = pyqtSignal()
def __init__(self, model, source, conf_threshold=0.5, iou_threshold=0.45):
super().__init__()
self.model = model
self.source = source # 可以是图片路径、视频路径或摄像头ID
self.conf_threshold = conf_threshold
self.iou_threshold = iou_threshold
self.is_running = True
def run(self):
"""线程主函数"""
try:
# 判断输入源类型
if isinstance(self.source, int) or str(self.source).endswith(('.mp4', '.avi', '.mov')):
self._process_video()
else:
self._process_image()
except Exception as e:
print(f"检测出错: {e}")
finally:
self.finished.emit()
def _process_video(self):
"""处理视频或摄像头流"""
cap = cv2.VideoCapture(self.source)
while self.is_running and cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# 保存原始帧用于显示
original = frame.copy()
# 用YOLO检测
results = self.model(
frame,
conf=self.conf_threshold,
iou=self.iou_threshold,
verbose=False # 不输出详细信息
)
# 绘制检测结果
annotated = results[0].plot()
# 提取检测信息
detections = []
for box in results[0].boxes:
class_id = int(box.cls[0])
class_name = self.model.names[class_id]
confidence = float(box.conf[0])
x, y, w, h = box.xywh[0].tolist()
detections.append({
'class': class_name,
'confidence': confidence,
'x': x,
'y': y,
'width': w,
'height': h
})
# 发送信号到主线程更新界面
self.frame_processed.emit(
cv2.cvtColor(original, cv2.COLOR_BGR2RGB),
cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB),
detections
)
# 控制帧率(大约30FPS)
self.msleep(33)
cap.release()
def _process_image(self):
"""处理单张图片"""
frame = cv2.imread(self.source)
if frame is None:
return
original = frame.copy()
results = self.model(
frame,
conf=self.conf_threshold,
iou=self.iou_threshold
)
annotated = results[0].plot()
detections = []
for box in results[0].boxes:
class_id = int(box.cls[0])
class_name = self.model.names[class_id]
confidence = float(box.conf[0])
x, y, w, h = box.xywh[0].tolist()
detections.append({
'class': class_name,
'confidence': confidence,
'x': x,
'y': y,
'width': w,
'height': h
})
self.frame_processed.emit(
cv2.cvtColor(original, cv2.COLOR_BGR2RGB),
cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB),
detections
)
def stop(self):
"""停止检测"""
self.is_running = False
class MainWindow(QMainWindow, Ui_MainWindow):
"""主窗口类,继承自动生成的界面"""
def __init__(self):
super().__init__()
self.setupUi(self) # 初始化界面
# 初始化变量
self.model = None
self.detection_thread = None
self.current_mode = None # 'image', 'video', 'camera'
self.video_writer = None
# 连接信号和槽
self._connect_signals()
# 加载模型
self._load_model()
# 设置窗口标题和大小
self.setWindowTitle("宠物品种识别系统")
self.resize(1200, 800)
def _connect_signals(self):
"""连接所有按钮和控件的信号"""
# 检测模式按钮
self.btn_image.clicked.connect(self.on_image_clicked)
self.btn_video.clicked.connect(self.on_video_clicked)
self.btn_camera.clicked.connect(self.on_camera_clicked)
self.btn_stop.clicked.connect(self.on_stop_clicked)
self.btn_save.clicked.connect(self.on_save_clicked)
# 参数调节
self.slider_confidence.valueChanged.connect(self.on_confidence_changed)
self.spinbox_confidence.valueChanged.connect(self.on_confidence_spinbox_changed)
self.slider_iou.valueChanged.connect(self.on_iou_changed)
self.spinbox_iou.valueChanged.connect(self.on_iou_spinbox_changed)
# 模型选择
self.combo_model.currentTextChanged.connect(self.on_model_changed)
def _load_model(self):
"""加载YOLO模型"""
try:
model_name = self.combo_model.currentText()
self.model = YOLO(f'{model_name}.pt')
self.statusbar.showMessage(f"模型 {model_name} 加载成功", 3000)
except Exception as e:
QMessageBox.critical(self, "错误", f"模型加载失败: {str(e)}")
def on_image_clicked(self):
"""图片检测按钮点击事件"""
if self.detection_thread and self.detection_thread.isRunning():
QMessageBox.warning(self, "提示", "请先停止当前检测任务")
return
# 打开文件对话框选择图片
file_path, _ = QFileDialog.getOpenFileName(
self,
"选择图片",
"",
"图片文件 (*.jpg *.jpeg *.png *.bmp)"
)
if file_path:
self.current_mode = 'image'
self._start_detection(file_path)
def on_video_clicked(self):
"""视频检测按钮点击事件"""
if self.detection_thread and self.detection_thread.isRunning():
QMessageBox.warning(self, "提示", "请先停止当前检测任务")
return
file_path, _ = QFileDialog.getOpenFileName(
self,
"选择视频",
"",
"视频文件 (*.mp4 *.avi *.mov)"
)
if file_path:
self.current_mode = 'video'
self._setup_video_writer(file_path)
self._start_detection(file_path)
def on_camera_clicked(self):
"""摄像头检测按钮点击事件"""
if self.detection_thread and self.detection_thread.isRunning():
QMessageBox.warning(self, "提示", "请先停止当前检测任务")
return
self.current_mode = 'camera'
self._start_detection(0) # 0表示默认摄像头
def _start_detection(self, source):
"""启动检测线程"""
# 清空之前的结果
self._clear_results()
# 获取当前参数
conf = self.spinbox_confidence.value()
iou = self.spinbox_iou.value()
# 创建并启动检测线程
self.detection_thread = DetectionThread(
model=self.model,
source=source,
conf_threshold=conf,
iou_threshold=iou
)
# 连接线程信号
self.detection_thread.frame_processed.connect(self.on_frame_processed)
self.detection_thread.finished.connect(self.on_detection_finished)
# 启动线程
self.detection_thread.start()
# 更新状态
mode_text = {
'image': '图片检测',
'video': '视频检测',
'camera': '摄像头检测'
}
self.statusbar.showMessage(f"正在执行{mode_text.get(self.current_mode, '')}...")
def on_frame_processed(self, original_frame, result_frame, detections):
"""接收到处理完的帧时调用"""
# 显示图片
self._display_image(self.label_original, original_frame)
self._display_image(self.label_result, result_frame)
# 更新结果表格
self._update_results_table(detections)
# 如果是视频模式,保存结果帧
if self.current_mode == 'video' and self.video_writer:
# 需要将RGB转回BGR才能用OpenCV保存
frame_bgr = cv2.cvtColor(result_frame, cv2.COLOR_RGB2BGR)
self.video_writer.write(frame_bgr)
def _display_image(self, label, image):
"""在QLabel上显示图片"""
h, w, ch = image.shape
bytes_per_line = ch * w
# 创建QImage
q_img = QImage(
image.data,
w, h,
bytes_per_line,
QImage.Format_RGB888
)
# 转换为QPixmap并显示
pixmap = QPixmap.fromImage(q_img)
# 缩放以适应Label大小,保持宽高比
scaled_pixmap = pixmap.scaled(
label.size(),
Qt.KeepAspectRatio,
Qt.SmoothTransformation
)
label.setPixmap(scaled_pixmap)
def _update_results_table(self, detections):
"""更新检测结果表格"""
# 清空表格
self.table_results.setRowCount(0)
# 添加新行
for i, det in enumerate(detections):
row_position = self.table_results.rowCount()
self.table_results.insertRow(row_position)
# 填充数据
self.table_results.setItem(row_position, 0, QTableWidgetItem(det['class']))
self.table_results.setItem(row_position, 1, QTableWidgetItem(f"{det['confidence']:.3f}"))
self.table_results.setItem(row_position, 2, QTableWidgetItem(f"{det['x']:.1f}"))
self.table_results.setItem(row_position, 3, QTableWidgetItem(f"{det['y']:.1f}"))
self.table_results.setItem(row_position, 4, QTableWidgetItem(f"{det['width']:.1f}"))
self.table_results.setItem(row_position, 5, QTableWidgetItem(f"{det['height']:.1f}"))
def _clear_results(self):
"""清空显示结果"""
self.label_original.clear()
self.label_result.clear()
self.table_results.setRowCount(0)
def _setup_video_writer(self, video_path):
"""设置视频写入器,用于保存检测结果"""
# 读取原视频信息
cap = cv2.VideoCapture(video_path)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
cap.release()
# 创建保存目录
import os
os.makedirs('results', exist_ok=True)
# 生成保存路径
import time
timestamp = time.strftime("%Y%m%d_%H%M%S")
save_path = f"results/detection_{timestamp}.mp4"
# 创建VideoWriter
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
self.video_writer = cv2.VideoWriter(save_path, fourcc, fps, (width, height))
def on_stop_clicked(self):
"""停止检测"""
if self.detection_thread:
self.detection_thread.stop()
self.detection_thread.wait()
if self.video_writer:
self.video_writer.release()
self.video_writer = None
self.statusbar.showMessage("检测已停止", 3000)
def on_save_clicked(self):
"""保存当前检测结果"""
if not hasattr(self, 'last_result_frame'):
QMessageBox.warning(self, "提示", "没有可保存的检测结果")
return
file_path, _ = QFileDialog.getSaveFileName(
self,
"保存图片",
"",
"PNG图片 (*.png);;JPEG图片 (*.jpg *.jpeg)"
)
if file_path:
# 将RGB转BGR保存
import cv2
bgr_frame = cv2.cvtColor(self.last_result_frame, cv2.COLOR_RGB2BGR)
cv2.imwrite(file_path, bgr_frame)
self.statusbar.showMessage(f"图片已保存到: {file_path}", 3000)
def on_confidence_changed(self, value):
"""置信度滑块值改变"""
conf = value / 100.0
self.spinbox_confidence.setValue(conf)
def on_confidence_spinbox_changed(self, value):
"""置信度数值框改变"""
self.slider_confidence.setValue(int(value * 100))
def on_iou_changed(self, value):
"""IoU滑块值改变"""
iou = value / 100.0
self.spinbox_iou.setValue(iou)
def on_iou_spinbox_changed(self, value):
"""IoU数值框改变"""
self.slider_iou.setValue(int(value * 100))
def on_model_changed(self, model_name):
"""切换模型"""
try:
self.model = YOLO(f'{model_name}.pt')
self.statusbar.showMessage(f"已切换到模型: {model_name}", 3000)
except Exception as e:
QMessageBox.critical(self, "错误", f"模型切换失败: {str(e)}")
def on_detection_finished(self):
"""检测线程完成时调用"""
if self.video_writer:
self.video_writer.release()
self.video_writer = None
self.statusbar.showMessage("检测完成", 3000)
def closeEvent(self, event):
"""窗口关闭事件"""
# 确保停止所有线程
self.on_stop_clicked()
event.accept()
def main():
app = QApplication(sys.argv)
window = MainWindow()
window.show()
sys.exit(app.exec_())
if __name__ == '__main__':
main()
```
这个主窗口类有近500行代码,但结构很清晰。我把它分成几个部分:
1. **DetectionThread类**:在后台运行检测任务,通过信号与主线程通信
2. **MainWindow类**:主界面,处理所有用户交互
3. **各种事件处理函数**:按钮点击、参数调节、模型切换等
运行`python main_window.py`,就能看到完整的宠物识别应用了。
## 4. 高级功能与优化技巧
基础功能有了,但要让应用更好用,还需要一些高级功能。这里分享几个我在实际开发中觉得很有用的技巧。
### 4.1 添加登录和用户管理
虽然是个桌面应用,但加上简单的用户系统可以让体验更完整。比如不同用户可以有各自的检测历史、收藏的品种等。
```python
# auth.py
import json
import hashlib
import os
from PyQt5.QtWidgets import QDialog, QVBoxLayout, QLabel, QLineEdit, QPushButton, QMessageBox
class LoginDialog(QDialog):
"""登录对话框"""
def __init__(self):
super().__init__()
self.setWindowTitle("用户登录")
self.setFixedSize(300, 200)
# 加载用户数据
self.users_file = 'users.json'
self.users = self._load_users()
# 创建界面
layout = QVBoxLayout()
layout.addWidget(QLabel("用户名:"))
self.input_username = QLineEdit()
layout.addWidget(self.input_username)
layout.addWidget(QLabel("密码:"))
self.input_password = QLineEdit()
self.input_password.setEchoMode(QLineEdit.Password)
layout.addWidget(self.input_password)
self.btn_login = QPushButton("登录")
self.btn_login.clicked.connect(self.on_login)
layout.addWidget(self.btn_login)
self.btn_register = QPushButton("注册新用户")
self.btn_register.clicked.connect(self.on_register)
layout.addWidget(self.btn_register)
self.setLayout(layout)
def _load_users(self):
"""加载用户数据"""
if os.path.exists(self.users_file):
with open(self.users_file, 'r', encoding='utf-8') as f:
return json.load(f)
return {}
def _save_users(self):
"""保存用户数据"""
with open(self.users_file, 'w', encoding='utf-8') as f:
json.dump(self.users, f, ensure_ascii=False, indent=2)
def _hash_password(self, password):
"""密码哈希(简单示例,实际应用需要加盐)"""
return hashlib.sha256(password.encode()).hexdigest()
def on_login(self):
"""登录按钮点击"""
username = self.input_username.text().strip()
password = self.input_password.text().strip()
if not username or not password:
QMessageBox.warning(self, "错误", "用户名和密码不能为空")
return
hashed_password = self._hash_password(password)
if username in self.users and self.users[username] == hashed_password:
self.current_user = username
self.accept() # 关闭对话框并返回Accepted
else:
QMessageBox.warning(self, "错误", "用户名或密码错误")
def on_register(self):
"""注册新用户"""
from PyQt5.QtWidgets import QInputDialog
username, ok = QInputDialog.getText(self, "注册", "请输入用户名:")
if not ok or not username:
return
if username in self.users:
QMessageBox.warning(self, "错误", "用户名已存在")
return
password, ok = QInputDialog.getText(
self, "注册", "请输入密码:",
QLineEdit.Password
)
if not ok or not password:
return
if len(password) < 6:
QMessageBox.warning(self, "错误", "密码至少需要6位")
return
# 确认密码
password2, ok = QInputDialog.getText(
self, "确认密码", "请再次输入密码:",
QLineEdit.Password
)
if not ok or password != password2:
QMessageBox.warning(self, "错误", "两次输入的密码不一致")
return
# 保存用户
self.users[username] = self._hash_password(password)
self._save_users()
QMessageBox.information(self, "成功", "注册成功!请登录")
```
然后在主程序启动时先显示登录对话框:
```python
# main.py
import sys
from PyQt5.QtWidgets import QApplication
from auth import LoginDialog
from main_window import MainWindow
def main():
app = QApplication(sys.argv)
# 显示登录对话框
login_dialog = LoginDialog()
if login_dialog.exec_() == LoginDialog.Accepted:
# 登录成功,显示主窗口
current_user = login_dialog.current_user
window = MainWindow(current_user)
window.show()
sys.exit(app.exec_())
else:
# 登录取消或失败,退出程序
sys.exit(0)
if __name__ == '__main__':
main()
```
### 4.2 实现检测历史记录
用户可能想查看之前的检测记录,我们可以添加这个功能:
```python
# history.py
import sqlite3
import json
from datetime import datetime
from PyQt5.QtWidgets import QWidget, QVBoxLayout, QTableWidget, QTableWidgetItem, QHeaderView
class DetectionHistory:
"""检测历史管理类"""
def __init__(self, db_path='detection_history.db'):
self.conn = sqlite3.connect(db_path)
self._create_table()
def _create_table(self):
"""创建历史记录表"""
cursor = self.conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS detection_history (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT NOT NULL,
detection_type TEXT NOT NULL, -- image/video/camera
source_path TEXT,
detection_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
results TEXT -- JSON格式的检测结果
)
''')
self.conn.commit()
def add_record(self, username, detection_type, source_path, results):
"""添加检测记录"""
cursor = self.conn.cursor()
cursor.execute('''
INSERT INTO detection_history
(username, detection_type, source_path, results)
VALUES (?, ?, ?, ?)
''', (username, detection_type, source_path, json.dumps(results)))
self.conn.commit()
return cursor.lastrowid
def get_user_history(self, username, limit=50):
"""获取用户的历史记录"""
cursor = self.conn.cursor()
cursor.execute('''
SELECT * FROM detection_history
WHERE username = ?
ORDER BY detection_time DESC
LIMIT ?
''', (username, limit))
records = []
for row in cursor.fetchall():
records.append({
'id': row[0],
'username': row[1],
'type': row[2],
'source': row[3],
'time': row[4],
'results': json.loads(row[5])
})
return records
def close(self):
"""关闭数据库连接"""
self.conn.close()
class HistoryWindow(QWidget):
"""历史记录查看窗口"""
def __init__(self, username, history_db):
super().__init__()
self.username = username
self.history_db = history_db
self.setWindowTitle(f"{username}的检测历史")
self.resize(800, 600)
# 创建界面
layout = QVBoxLayout()
self.table = QTableWidget()
self.table.setColumnCount(5)
self.table.setHorizontalHeaderLabels([
'时间', '检测类型', '来源', '检测数量', '主要品种'
])
# 设置表格属性
self.table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
self.table.setAlternatingRowColors(True)
layout.addWidget(self.table)
self.setLayout(layout)
# 加载数据
self.load_history()
def load_history(self):
"""加载历史记录"""
records = self.history_db.get_user_history(self.username)
self.table.setRowCount(len(records))
for row_idx, record in enumerate(records):
# 解析结果
results = record['results']
detection_count = len(results) if isinstance(results, list) else 0
# 找出置信度最高的品种
main_breed = "无"
if detection_count > 0:
# 假设results是检测结果列表
highest_conf = max(results, key=lambda x: x.get('confidence', 0))
main_breed = highest_conf.get('class', '未知')
# 填充表格
self.table.setItem(row_idx, 0, QTableWidgetItem(record['time']))
self.table.setItem(row_idx, 1, QTableWidgetItem(record['type']))
self.table.setItem(row_idx, 2, QTableWidgetItem(record['source'] or '摄像头'))
self.table.setItem(row_idx, 3, QTableWidgetItem(str(detection_count)))
self.table.setItem(row_idx, 4, QTableWidgetItem(main_breed))
```
在主窗口中添加一个查看历史的按钮:
```python
# 在MainWindow类的__init__中添加
self.btn_history = QPushButton("查看历史")
self.btn_history.clicked.connect(self.on_view_history)
# 添加事件处理函数
def on_view_history(self):
"""查看历史记录"""
from history import HistoryWindow
self.history_window = HistoryWindow(self.current_user, self.history_db)
self.history_window.show()
```
### 4.3 模型性能优化技巧
当应用运行在资源有限的设备上时,这些优化技巧很有用:
**1. 模型量化(减小模型大小)**
```python
def quantize_model(model_path, output_path):
"""量化模型,减小文件大小"""
import torch
from ultralytics import YOLO
# 加载模型
model = YOLO(model_path)
# 转换为TorchScript
model.export(format='torchscript')
# 量化(需要PyTorch 1.3+)
quantized_model = torch.quantization.quantize_dynamic(
model.model,
{torch.nn.Linear, torch.nn.Conv2d},
dtype=torch.qint8
)
# 保存量化模型
torch.jit.save(torch.jit.script(quantized_model), output_path)
print(f"量化模型已保存到: {output_path}")
```
**2. 使用ONNX Runtime加速推理**
```python
def convert_to_onnx(model_path, output_path):
"""转换为ONNX格式,可以用ONNX Runtime加速"""
from ultralytics import YOLO
model = YOLO(model_path)
# 导出为ONNX
success = model.export(
format='onnx',
imgsz=640,
opset=12,
simplify=True,
dynamic=False # 固定输入尺寸,速度更快
)
if success:
print(f"ONNX模型已保存到: {output_path}")
# 使用ONNX Runtime推理
import onnxruntime as ort
# 创建推理会话
session = ort.InferenceSession(output_path)
# 准备输入
import numpy as np
dummy_input = np.random.randn(1, 3, 640, 640).astype(np.float32)
# 推理
outputs = session.run(None, {'images': dummy_input})
print("ONNX推理成功!")
```
**3. 多尺度推理提升小目标检测**
```python
def multi_scale_inference(model, image, scales=[0.5, 1.0, 1.5]):
"""多尺度推理,提升小目标检测效果"""
import cv2
import numpy as np
all_results = []
for scale in scales:
# 缩放图片
h, w = image.shape[:2]
new_w, new_h = int(w * scale), int(h * scale)
resized = cv2.resize(image, (new_w, new_h))
# 推理
results = model(resized, conf=0.3, iou=0.4)
# 将检测框缩放回原图尺寸
for result in results:
for box in result.boxes:
# 缩放坐标
x1, y1, x2, y2 = box.xyxy[0].tolist()
x1, x2 = x1 / scale, x2 / scale
y1, y2 = y1 / scale, y2 / scale
all_results.append({
'box': [x1, y1, x2, y2],
'conf': float(box.conf[0]),
'cls': int(box.cls[0])
})
# 非极大值抑制,去除重复框
return non_max_suppression(all_results)
def non_max_suppression(detections, iou_threshold=0.5):
"""非极大值抑制"""
if not detections:
return []
# 按置信度排序
detections.sort(key=lambda x: x['conf'], reverse=True)
keep = []
while detections:
# 取置信度最高的
best = detections.pop(0)
keep.append(best)
# 计算与剩余框的IoU
to_remove = []
for i, det in enumerate(detections):
iou = calculate_iou(best['box'], det['box'])
if iou > iou_threshold:
to_remove.append(i)
# 移除重叠框
for idx in reversed(to_remove):
detections.pop(idx)
return keep
```
### 4.4 界面美化与用户体验
好的界面能让应用更受欢迎。这里分享几个PyQt5的美化技巧:
**1. 自定义样式表**
```python
def apply_dark_theme(window):
"""应用深色主题"""
dark_stylesheet = """
QMainWindow {
background-color: #1e1e1e;
}
QLabel {
color: #ffffff;
font-size: 12px;
}
QPushButton {
background-color: #2d2d30;
color: #ffffff;
border: 1px solid #3e3e42;
border-radius: 4px;
padding: 6px 12px;
font-weight: bold;
}
QPushButton:hover {
background-color: #3e3e42;
border-color: #007acc;
}
QPushButton:pressed {
background-color: #0e639c;
}
QTableWidget {
background-color: #252526;
color: #cccccc;
gridline-color: #3e3e42;
border: 1px solid #3e3e42;
}
QTableWidget::item {
padding: 4px;
}
QTableWidget::item:selected {
background-color: #094771;
}
QHeaderView::section {
background-color: #2d2d30;
color: #ffffff;
padding: 6px;
border: 1px solid #3e3e42;
}
QSlider::groove:horizontal {
height: 6px;
background: #3e3e42;
border-radius: 3px;
}
QSlider::handle:horizontal {
background: #007acc;
width: 16px;
height: 16px;
margin: -5px 0;
border-radius: 8px;
}
QComboBox {
background-color: #2d2d30;
color: #ffffff;
border: 1px solid #3e3e42;
border-radius: 4px;
padding: 4px;
}
QComboBox::drop-down {
border: none;
}
QComboBox QAbstractItemView {
background-color: #2d2d30;
color: #ffffff;
selection-background-color: #094771;
}
"""
window.setStyleSheet(dark_stylesheet)
```
**2. 添加动画效果**
```python
from PyQt5.QtCore import QPropertyAnimation, QEasingCurve
def add_button_animation(button):
"""给按钮添加悬停动画"""
# 创建动画对象
animation = QPropertyAnimation(button, b"geometry")
animation.setDuration(200) # 200毫秒
animation.setEasingCurve(QEasingCurve.OutBack)
# 保存原始位置
original_geometry = button.geometry()
# 鼠标进入事件
def on_enter(event):
# 稍微放大按钮
new_geo = original_geometry.adjusted(-2, -2, 2, 2)
animation.setEndValue(new_geo)
animation.start()
# 鼠标离开事件
def on_leave(event):
animation.setEndValue(original_geometry)
animation.start()
button.enterEvent = on_enter
button.leaveEvent = on_leave
```
**3. 实现拖放功能**
```python
from PyQt5.QtCore import Qt
class DragDropLabel(QLabel):
"""支持拖放图片的Label"""
def __init__(self, parent=None):
super().__init__(parent)
self.setAcceptDrops(True)
self.setText("拖放图片到这里")
self.setAlignment(Qt.AlignCenter)
self.setStyleSheet("""
border: 2px dashed #aaa;
border-radius: 8px;
padding: 20px;
color: #888;
""")
def dragEnterEvent(self, event):
"""拖拽进入事件"""
if event.mimeData().hasUrls():
event.acceptProposedAction()
def dropEvent(self, event):
"""放下事件"""
urls = event.mimeData().urls()
if urls:
file_path = urls[0].toLocalFile()
if file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
# 加载并显示图片
pixmap = QPixmap(file_path)
if not pixmap.isNull():
self.setPixmap(
pixmap.scaled(
self.size(),
Qt.KeepAspectRatio,
Qt.SmoothTransformation
)
)
# 发出信号,通知主窗口
self.parent().image_dropped.emit(file_path)
```
把这些功能整合起来,你的宠物识别应用就会变得既强大又好用。从数据准备到模型训练,从界面开发到功能优化,整个过程虽然有些复杂,但每一步都有明确的实现方法。最重要的是,你完全掌控了整个系统,可以根据自己的需求随时调整和扩展。