基于改进 YOLOv8 的美食图像分割系统实战与源码解析

实时检测与可视化

这是核心的推理逻辑部分。我们使用了 OpenCV 读取摄像头流，结合深度学习模型进行预测。为了让结果更直观，代码中实现了中文文本绘制和动态颜色生成。注意 generate_color_based_on_name 函数利用 MD5 哈希确保同一类别始终显示相同颜色，方便追踪。

import cv2
import numpy as np
from PIL import ImageFont, ImageDraw, Image
from hashlib import md5
from model import Web_Detector
from chinese_name_list import Label_list

def generate_color_based_on_name(name):
    hash_object = md5(name.encode())
    hex_color = hash_object.hexdigest()[:6]
    r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
    return (b, g, r)  # OpenCV 使用 BGR 格式

def draw_with_chinese(image, text, position, font_size=20, color=(255, 0, 0)):
    image_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(image_pil)
    font = ImageFont.truetype("simsun.ttc", font_size, encoding="unic")
    draw.text(position, text, font=font, fill=color)
    return cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)

def draw_detections(image, info):
    name, bbox = info['class_name'], info['bbox']
    x1, y1, x2, y2 = bbox
    cv2.rectangle(image, (x1, y1), (x2, y2), color=(0, 0, 255), thickness=3)
    image = draw_with_chinese(image, name, (x1, y1 - 10), font_size=20)
    return image

def process_frame(model, image):
    pre_img = model.preprocess(image)
    pred = model.predict(pre_img)
    det = pred[0]
    if det is not None and len(det):
        det_info = model.postprocess(pred)
        for info in det_info:
            image = draw_detections(image, info)
    return image

if __name__ == "__main__":
    model = Web_Detector()
    model.load_model("./weights/yolov8s-seg.pt")
    cap = cv2.VideoCapture(0)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret: break
        processed_frame = process_frame(model, frame)
        cv2.imshow('Camera Feed', processed_frame)
        if cv2.waitKey(1) & 0xFF == ord('q'): break
    cap.release()
    cv2.destroyAllWindows()

模型接口定义

这里展示了 FastSAM 的类定义，它是 Ultralytics YOLO 框架中的一个模型接口，用于图像分割任务。初始化时会自动检查模型后缀，确保加载的是预训练权重而非配置文件。

from pathlib import Path
from ultralytics.engine.model import Model
from .predict import FastSAMPredictor
from .val import FastSAMValidator

class FastSAM(Model):
    """FastSAM 模型接口"""
    def __init__(self, model='FastSAM-x.pt'):
        if str(model) == 'FastSAM.pt':
            model = 'FastSAM-x.pt'
        assert Path(model).suffix not in ('.yaml', '.yml'), 'FastSAM models only support pre-trained models.'
        super().__init__(model=model, task='segment')

    @property
    def task_map(self):
        return {'segment': {'predictor': FastSAMPredictor, 'validator': FastSAMValidator}}

姿态估计训练器

虽然主任务是分割，但系统也集成了姿态估计功能。PoseTrainer 继承自 DetectionTrainer，专门处理关键点形状属性。注意其中对 Apple MPS 设备的警告，建议在使用 Pose 模型时使用 CPU 以避免已知 bug。

from ultralytics.models import yolo
from ultralytics.nn.tasks import PoseModel
from ultralytics.utils import DEFAULT_CFG, LOGGER

class PoseTrainer(yolo.detect.DetectionTrainer):
    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
        if overrides is None:
            overrides = {}
        overrides['task'] = 'pose'
        super().__init__(cfg, overrides, _callbacks)
        if isinstance(self.args.device, str) and self.args.device.lower() == 'mps':
            LOGGER.warning("WARNING ⚠️ Apple MPS known Pose bug. Recommend 'device=cpu'.")

    def get_model(self, cfg=None, weights=None, verbose=True):
        model = PoseModel(cfg, ch=3, nc=self.data['nc'], data_kpt_shape=self.data['kpt_shape'], verbose=verbose)
        if weights:
            model.load(weights)
        return model

训练脚本

这是实际触发训练的入口。脚本会动态调整 YAML 配置文件中的路径，确保相对于当前目录正确加载数据集。显存不足时记得调小 batch 参数。

import os
import torch
import yaml
from ultralytics import YOLO

def abs_path(path, path_type='current'):
    # 假设此处有实现绝对路径的逻辑
    return os.path.abspath(path)

if __name__ == '__main__':
    workers = 1
    batch = 8
    device = "0" if torch.cuda.is_available() else "cpu"
    
    data_path = abs_path(f'datasets/data/data.yaml')
    with open(data_path, 'r') as file:
        data = yaml.load(file, Loader=yaml.FullLoader)
    
    if 'train' in data and 'val' in data and 'test' in data:
        directory_path = os.path.dirname(data_path.replace(os.sep, '/'))
        data['train'] = directory_path + '/train'
        data['val'] = directory_path + '/val'
        data['test'] = directory_path + '/test'
        with open(data_path, 'w') as file:
            yaml.safe_dump(data, file, sort_keys=False)
    
    model = YOLO(r"./config/yolov8-seg-C2f-Faster.yaml").load("./weights/yolov8s-seg.pt")
    results = model.train(
        data=data_path,
        device=device,
        workers=workers,
        imgsz=640,
        epochs=100,
        batch=batch
    )

目标跟踪工具

这部分代码处理轨迹匹配问题。linear_assignment 函数利用成本矩阵进行最优分配，iou_distance 计算交并比作为成本。这些底层算法是保证多目标跟踪稳定性的关键。

import numpy as np
import scipy
from scipy.spatial.distance import cdist
from ultralytics.utils.metrics import bbox_ioa

def linear_assignment(cost_matrix, thresh, use_lap=True):
    if cost_matrix.size == 0:
        return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))
    if use_lap:
        try:
            import lap
            _, x, y = lap.lapjv(cost_matrix, extend_cost=True, cost_limit=thresh)
            matches = [[ix, mx] for ix, mx in enumerate(x) if mx >= 0]
            unmatched_a = np.where(x < 0)[0]
            unmatched_b = np.where(y < 0)[0]
        except ImportError:
            # Fallback or error handling
            pass
    else:
        x, y = scipy.optimize.linear_sum_assignment(cost_matrix)
        matches = np.asarray([[x[i], y[i]] for i in range(len(x)) if cost_matrix[x[i], y[i]] <= thresh])
    return matches, list(unmatched_a), list(unmatched_b)

def iou_distance(atracks, btracks):
    atlbrs = [track.tlbr for track in atracks] if atracks else []
    btlbrs = [track.tlbr for track in btracks] if btracks else []
    ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
    if len(atlbrs) and len(btlbrs):
        ious = bbox_ioa(np.ascontiguousarray(atlbrs, dtype=np.float32), 
                        np.ascontiguousarray(btlbrs, dtype=np.float32), iou=True)
    return 1 - ious

总结

整个系统从数据处理到模型训练，再到前端展示和实时推理，形成了一套完整的闭环。通过改进 YOLOv8 架构并结合自定义的后处理逻辑，我们在保持高精度的同时优化了响应速度。对于希望深入计算机视觉工程落地的开发者来说，这套代码提供了很好的参考范本。