基于YOLOv8的药物识别系统：从数据集构建到界面开发全流程

药物识别在医疗、药房管理和家庭用药安全中一直是个痛点。传统的人工识别方式效率低下且容易出错，而市面上的专业识别设备价格昂贵。有没有一种方法，能让普通开发者也能快速搭建一个高精度的药物识别系统？

YOLOv8作为当前最先进的目标检测算法之一，在速度和精度之间取得了很好的平衡。但很多教程只停留在理论层面，真正要把YOLOv8应用到药物识别这个具体场景，还需要解决数据集构建、模型训练、界面开发等一系列实际问题。

本文将从零开始，带你完整实现一个基于YOLOv8的药物识别检测系统。不仅仅是跑通Demo，更重要的是分享在实际项目中容易踩坑的细节：如何选择合适的数据集标注工具、怎样调整模型参数获得最佳效果、以及如何将训练好的模型封装成可用的UI界面。

1. 药物识别系统的实际价值与应用场景

药物识别系统看似小众，但在多个场景下都有重要价值。在医院药房，新入职的药师需要快速熟悉上千种药品的外观；在家庭用药场景，老人经常混淆相似包装的药物；在药品监管领域，需要快速识别假药和过期药品。

传统的解决方案要么依赖人工经验，要么需要昂贵的专业设备。而基于深度学习的识别系统，一旦训练完成，可以以极低的边际成本进行大规模部署。YOLOv8作为单阶段检测算法，特别适合这类需要实时识别的场景。

与传统的CNN分类模型相比，YOLOv8不仅能识别药物种类，还能精确标出药品在图像中的位置，这对于处理包含多个药物的复杂场景尤为重要。在实际应用中，用户可能同时拍摄多款药物，需要系统能够分别识别并标注。

2. YOLOv8的核心优势与药物识别适配性

YOLOv8在药物识别任务中表现出色的原因主要有三点：首先是检测精度高，能够准确区分外观相似的药物；其次是推理速度快，满足实时检测需求；最后是模型尺寸相对较小，便于部署到各种设备。

药物识别与其他物体检测任务相比有其特殊性：药物通常有标准化的外观，但不同厂家生产的同种药物包装可能差异很大；药物在图像中的尺度变化较大，有些是整盒拍摄，有些是单片特写；光照条件、拍摄角度等因素都会影响识别效果。

YOLOv8的Anchor-Free设计和先进的标签分配策略，使其在处理尺度变化大的物体时表现优异。同时，YOLOv8提供了n、s、m、l、x五种不同规模的模型，可以根据实际需求在精度和速度之间进行权衡。

3. 环境配置与依赖安装

在开始项目之前，需要确保环境配置正确。推荐使用Python 3.8或3.9版本，避免使用过新或过旧的Python版本可能带来的兼容性问题。

3.1 创建虚拟环境

首先创建独立的Python环境，避免包冲突：

conda create -n yolov8-drug python=3.9 conda activate yolov8-drug

3.2 安装核心依赖

pip install ultralytics==8.0.0 pip install opencv-python==4.8.0 pip install pillow==9.5.0 pip install matplotlib==3.7.0

3.3 验证安装

创建测试脚本验证环境是否正确：

# test_environment.py import ultralytics import cv2 import torch print(f"Ultralytics版本: {ultralytics.__version__}") print(f"OpenCV版本: {cv2.__version__}") print(f"PyTorch版本: {torch.__version__}") print(f"CUDA是否可用: {torch.cuda.is_available()}") # 测试YOLOv8模型加载 from ultralytics import YOLO model = YOLO('yolov8n.pt') print("YOLOv8模型加载成功!")

运行该脚本应该显示各库的版本信息，且不报错。

4. 药物数据集准备与标注

数据集的质量直接决定模型的识别效果。药物数据集可以从多个渠道获取：公开数据集、网络爬取、实际拍摄等。

4.1 数据收集策略

建议采用多源数据收集策略：

• 从公开医疗数据集中获取基础数据
• 通过爬虫获取电商平台的药物图片
• 实际拍摄不同光照条件下的药物照片
• 使用数据增强技术扩充数据集

4.2 数据标注工具选择

推荐使用LabelImg或CVAT进行标注：

# 安装LabelImg pip install labelimg labelimg # 启动标注工具

标注时需要注意以下几点：

• 标注框要紧密贴合药物边缘
• 同类药物使用相同标签名称
• 标注文件保存为YOLO格式（txt文件）

4.3 数据集目录结构

规范的数据集结构便于后续训练：

drug_dataset/ ├── images/ │ ├── train/ │ │ ├── drug_001.jpg │ │ └── drug_002.jpg │ └── val/ │ ├── drug_101.jpg │ └── drug_102.jpg └── labels/ ├── train/ │ ├── drug_001.txt │ └── drug_002.txt └── val/ ├── drug_101.txt └── drug_102.txt

4.4 数据集配置文件

创建数据集配置文件drug_dataset.yaml：

# drug_dataset.yaml path: /path/to/drug_dataset train: images/train val: images/val nc: 10 # 药物类别数量 names: ['aspirin', 'vitamin_c', 'antibiotic', 'pain_reliever', 'antihistamine', 'blood_pressure', 'diabetes', 'cholesterol', 'antacid', 'supplement']

5. YOLOv8模型训练全流程

5.1 基础训练配置

使用预训练权重进行迁移学习：

# train_basic.py from ultralytics import YOLO # 加载预训练模型 model = YOLO('yolov8n.pt') # 开始训练 results = model.train( data='drug_dataset.yaml', epochs=100, imgsz=640, batch=16, device=0, # 使用GPU workers=4, patience=10, # 早停耐心值 save=True, pretrained=True )

5.2 高级训练参数调优

针对药物识别任务调整超参数：

# train_advanced.py from ultralytics import YOLO model = YOLO('yolov8s.pt') results = model.train( data='drug_dataset.yaml', epochs=150, imgsz=640, batch=8, # 小批量适合复杂场景 lr0=0.01, # 初始学习率 lrf=0.01, # 最终学习率 momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, box=7.5, # 调整损失权重 cls=0.5, dfl=1.5, hsv_h=0.015, # 色相增强 hsv_s=0.7, # 饱和度增强 hsv_v=0.4, # 明度增强 degrees=0.0, # 旋转角度 translate=0.1, # 平移 scale=0.5, # 缩放 shear=0.0, # 剪切 perspective=0.0, # 透视变换 flipud=0.0, # 上下翻转 fliplr=0.5, # 左右翻转 mosaic=1.0, # 马赛克增强 mixup=0.0, # MixUp增强 copy_paste=0.0 # 复制粘贴增强 )

5.3 训练过程监控

实时监控训练指标，及时调整策略：

# monitor_training.py import matplotlib.pyplot as plt from ultralytics.utils.plots import plot_results # 绘制训练结果 results = model.train(...) plot_results('runs/detect/train/results.csv', save=True) # 实时监控关键指标 def analyze_training(results_path): import pandas as pd results_df = pd.read_csv(results_path) plt.figure(figsize=(12, 8)) # 损失函数变化 plt.subplot(2, 2, 1) plt.plot(results_df['epoch'], results_df['train/box_loss'], label='Box Loss') plt.plot(results_df['epoch'], results_df['train/cls_loss'], label='Cls Loss') plt.legend() plt.title('Training Loss') # 精度指标 plt.subplot(2, 2, 2) plt.plot(results_df['epoch'], results_df['metrics/precision(B)'], label='Precision') plt.plot(results_df['epoch'], results_df['metrics/recall(B)'], label='Recall') plt.legend() plt.title('Precision & Recall') plt.tight_layout() plt.savefig('training_analysis.png')

6. 模型评估与性能优化

6.1 评估模型性能

使用验证集评估训练好的模型：

# evaluate_model.py from ultralytics import YOLO # 加载最佳模型 model = YOLO('runs/detect/train/weights/best.pt') # 在验证集上评估 metrics = model.val( data='drug_dataset.yaml', split='val', batch=16, imgsz=640, conf=0.25, # 置信度阈值 iou=0.45, # IoU阈值 device=0 ) print(f"mAP50: {metrics.box.map50}") print(f"mAP50-95: {metrics.box.map}") print(f"Precision: {metrics.box.mp}") print(f"Recall: {metrics.box.mr}")

6.2 混淆矩阵分析

# confusion_matrix.py from ultralytics.utils.plots import plot_confusion_matrix import matplotlib.pyplot as plt # 绘制混淆矩阵 plot_confusion_matrix( matrix=metrics.confusion_matrix, normalize=True, save_dir='runs/detect/train/', names=model.names )

6.3 性能优化策略

针对识别效果不佳的类别进行优化：

# optimize_performance.py def analyze_class_performance(metrics, class_names): """分析各类别性能""" precision_per_class = metrics.box.p_per_class recall_per_class = metrics.box.r_per_class performance_data = [] for i, name in enumerate(class_names): performance_data.append({ 'class': name, 'precision': precision_per_class[i], 'recall': recall_per_class[i], 'f1_score': 2 * (precision_per_class[i] * recall_per_class[i]) / (precision_per_class[i] + recall_per_class[i] + 1e-16) }) # 按F1分数排序 performance_data.sort(key=lambda x: x['f1_score']) print("性能最差的类别:") for data in performance_data[:3]: print(f"{data['class']}: Precision={data['precision']:.3f}, " f"Recall={data['recall']:.3f}, F1={data['f1_score']:.3f}") return performance_data

7. 图形界面开发与集成

7.1 使用Gradio构建Web界面

Gradio适合快速构建演示界面：

# app_gradio.py import gradio as gr import cv2 import numpy as np from ultralytics import YOLO # 加载训练好的模型 model = YOLO('runs/detect/train/weights/best.pt') def predict_drugs(image): """药物识别预测函数""" # 运行推理 results = model(image, conf=0.25) # 绘制检测结果 annotated_image = results[0].plot() # 提取检测信息 detections = [] for result in results: boxes = result.boxes for box in boxes: cls_id = int(box.cls[0]) confidence = float(box.conf[0]) class_name = model.names[cls_id] detections.append({ 'class': class_name, 'confidence': confidence, 'bbox': box.xyxy[0].tolist() }) return annotated_image, detections # 创建界面 iface = gr.Interface( fn=predict_drugs, inputs=gr.Image(label="上传药物图片"), outputs=[ gr.Image(label="检测结果"), gr.JSON(label="检测详情") ], title="YOLOv8药物识别系统", description="上传药物图片，系统将自动识别药物类型并标注位置" ) iface.launch(server_name="0.0.0.0", server_port=7860)

7.2 使用Streamlit构建专业界面

Streamlit适合构建更复杂的数据应用：

# app_streamlit.py import streamlit as st import cv2 import numpy as np from PIL import Image from ultralytics import YOLO # 页面配置 st.set_page_config( page_title="药物识别系统", page_icon="💊", layout="wide" ) # 标题和说明 st.title("💊 YOLOv8药物识别检测系统") st.markdown("上传药物图片，系统将自动识别药物类型并标注检测框") # 侧边栏配置 st.sidebar.header("模型配置") confidence_threshold = st.sidebar.slider("置信度阈值", 0.1, 0.9, 0.25) iou_threshold = st.sidebar.slider("IoU阈值", 0.1, 0.9, 0.45) # 文件上传 uploaded_file = st.file_uploader( "选择药物图片", type=['jpg', 'jpeg', 'png'] ) @st.cache_resource def load_model(): """加载模型（缓存）""" return YOLO('runs/detect/train/weights/best.pt') model = load_model() if uploaded_file is not None: # 读取图片 image = Image.open(uploaded_file) image_np = np.array(image) col1, col2 = st.columns(2) with col1: st.subheader("原图") st.image(image, use_column_width=True) with col2: st.subheader("检测结果") # 运行推理 with st.spinner("正在识别药物..."): results = model( image_np, conf=confidence_threshold, iou=iou_threshold ) # 绘制结果 annotated_image = results[0].plot() st.image(annotated_image, use_column_width=True) # 显示检测详情 st.subheader("检测详情") detections = [] for result in results: boxes = result.boxes for box in boxes: cls_id = int(box.cls[0]) confidence = float(box.conf[0]) class_name = model.names[cls_id] bbox = box.xyxy[0].tolist() detections.append({ '药物类型': class_name, '置信度': f"{confidence:.3f}", '位置': [f"{coord:.1f}" for coord in bbox] }) if detections: st.table(detections) else: st.warning("未检测到药物")

8. 系统部署与性能优化

8.1 模型导出与优化

将模型导出为不同格式以适应不同部署场景：

# export_models.py from ultralytics import YOLO model = YOLO('runs/detect/train/weights/best.pt') # 导出为ONNX格式（适合跨平台部署） model.export(format='onnx', imgsz=640, simplify=True) # 导出为TensorRT格式（适合GPU加速） model.export(format='engine', imgsz=640, device=0) # 导出为OpenVINO格式（适合Intel硬件） model.export(format='openvino', imgsz=640) print("模型导出完成!")

8.2 使用FastAPI构建API服务

# api_fastapi.py from fastapi import FastAPI, File, UploadFile from fastapi.responses import JSONResponse import cv2 import numpy as np from ultralytics import YOLO import io from PIL import Image app = FastAPI(title="药物识别API") # 加载模型 model = YOLO('runs/detect/train/weights/best.pt') @app.post("/predict") async def predict_drugs(file: UploadFile = File(...)): """药物识别API接口""" # 读取图片 image_data = await file.read() image = Image.open(io.BytesIO(image_data)) image_np = np.array(image) # 运行推理 results = model(image_np) # 处理结果 detections = [] for result in results: boxes = result.boxes for box in boxes: cls_id = int(box.cls[0]) confidence = float(box.conf[0]) class_name = model.names[cls_id] bbox = box.xyxy[0].tolist() detections.append({ 'class': class_name, 'confidence': confidence, 'bbox': bbox }) return JSONResponse({ 'status': 'success', 'detections': detections, 'count': len(detections) }) @app.get("/health") async def health_check(): """健康检查接口""" return {"status": "healthy"} if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)

9. 实际应用中的挑战与解决方案

9.1 光照条件变化应对

药物识别在实际应用中最大的挑战之一是光照条件的变化：

# illumination_handling.py import cv2 import numpy as np def preprocess_image(image): """图像预处理，增强光照鲁棒性""" # 转换为HSV色彩空间 hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) # 直方图均衡化（亮度通道） hsv[:,:,2] = cv2.equalizeHist(hsv[:,:,2]) # 转换回BGR enhanced = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) # 伽马校正 gamma = 1.5 inv_gamma = 1.0 / gamma table = np.array([((i / 255.0) ** inv_gamma) * 255 for i in np.arange(0, 256)]).astype("uint8") gamma_corrected = cv2.LUT(enhanced, table) return gamma_corrected

9.2 多尺度检测优化

药物在图像中的尺度变化很大，需要优化多尺度检测：

# multi_scale_detection.py def multi_scale_predict(model, image, scales=[0.5, 1.0, 1.5]): """多尺度预测融合""" all_results = [] for scale in scales: # 调整图像尺寸 new_width = int(image.shape[1] * scale) new_height = int(image.shape[0] * scale) resized_image = cv2.resize(image, (new_width, new_height)) # 预测 results = model(resized_image) # 调整边界框坐标 for result in results: boxes = result.boxes for box in boxes: # 缩放回原图尺寸 bbox = box.xyxy[0].cpu().numpy() / scale box.xyxy[0] = torch.tensor(bbox) all_results.append(box) # 非极大值抑制融合结果 if all_results: # 实现NMS逻辑... pass return all_results

10. 系统测试与验证

10.1 创建测试用例

# test_system.py import unittest import cv2 import numpy as np from ultralytics import YOLO class TestDrugRecognition(unittest.TestCase): def setUp(self): self.model = YOLO('runs/detect/train/weights/best.pt') # 创建测试图像 self.test_image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8) def test_model_loading(self): """测试模型加载""" self.assertIsNotNone(self.model) self.assertTrue(hasattr(self.model, 'names')) def test_prediction_output(self): """测试预测输出格式""" results = self.model(self.test_image) self.assertTrue(len(results) > 0) # 检查结果结构 result = results[0] self.assertTrue(hasattr(result, 'boxes')) def test_confidence_threshold(self): """测试置信度阈值""" results = self.model(self.test_image, conf=0.5) # 验证所有检测的置信度都大于阈值 if len(results[0].boxes) > 0: for box in results[0].boxes: self.assertGreaterEqual(float(box.conf[0]), 0.5) if __name__ == '__main__': unittest.main()

10.2 性能基准测试

# benchmark.py import time import cv2 import numpy as np from ultralytics import YOLO def benchmark_model(model, image_size=(640, 640), num_tests=100): """模型性能基准测试""" # 创建测试图像 test_image = np.random.randint(0, 255, (*image_size, 3), dtype=np.uint8) # 预热 _ = model(test_image) # 性能测试 start_time = time.time() for _ in range(num_tests): _ = model(test_image) end_time = time.time() avg_inference_time = (end_time - start_time) / num_tests fps = 1.0 / avg_inference_time print(f"平均推理时间: {avg_inference_time*1000:.2f}ms") print(f"FPS: {fps:.2f}") return avg_inference_time, fps # 运行基准测试 model = YOLO('runs/detect/train/weights/best.pt') benchmark_model(model)

药物识别系统的实际部署需要考虑更多工程细节，比如错误处理、日志记录、监控告警等。建议在生产环境中添加完整的异常处理机制，确保系统的稳定性和可靠性。

对于想要进一步优化的开发者，可以考虑使用模型蒸馏技术减小模型尺寸，或者集成多个模型进行集成学习提高识别准确率。在实际项目中，持续收集用户反馈和数据，定期重新训练模型，才能让系统保持最佳性能。