PySpark + YOLOv8 分布式图像目标检测实战｜大数据与 AI 融合教学项目

前言

传统单机图像推理在面对海量数据时，串行处理带来的时间开销往往难以承受。本文将 PySpark 分布式计算框架 与 YOLOv8 目标检测算法 深度结合，构建一套通用的分布式图像推理流水线。项目涵盖数据集格式转换、模型训练与分布式批量推理三个核心环节，无特定业务场景绑定，可直接迁移至任意图像检测任务，适合大数据与计算机视觉交叉方向的学习与实践。

项目核心学习价值

1. 掌握大数据框架与深度学习模型融合的标准开发流程；
2. 深入理解 mapPartitionsWithIndex 分区算子的优势：每个分区仅加载一次模型，避免重复初始化造成的内存/显存浪费；
3. 熟练完成 VOC 标注到 YOLO 格式的转换，以及数据集的分层划分；
4. 实现分布式批量图像推理、检测框可视化与全局目标数量统计；
5. 提供独立可运行的完整源码，支持 Linux (Ubuntu) + Anaconda 环境下一键复现。

---

一、实验环境准备

1. 基础要求

推荐操作系统：Ubuntu 20.04+。需预先部署 Apache Spark 与 Anaconda。

2. 创建并激活虚拟环境

conda create -n spark-yolo python=3.8
conda activate spark-yolo

3. 安装依赖库

pip install pyspark==3.5.8
pip install ultralytics==8.4.60
pip install scikit-learn==1.3.2
pip install opencv-python numpy

4. 项目目录结构

spark-yolo-project/
├── dataset/                    # 原始 VOC 数据集
│   ├── images/                 # 原始图像
│   └── annotations/            # VOC XML 标注
├── convert_and_split_dataset.py  # 格式转换与数据集划分
├── train_yolo.py               # YOLOv8 模型训练
├── dist_predict.py             # PySpark 分布式推理
└── yolov8n.pt                  # 预训练权重

---

二、步骤 1：VOC 数据集转换与分层划分

业务说明

VOC 格式（图像 + XML）是常见的标注标准，Ultralytics YOLO 只接受归一化 TXT 标签。因此需完成坐标格式转换，并按 6:2:2 比例分层拆分为训练集、验证集、测试集，以保持数据分布的一致性。

完整脚本 convert_and_split_dataset.py

import os
import shutil
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import List, Tuple
from sklearn.model_selection import train_test_split

# ======================== 可配置参数 ========================
DATASET_ROOT: str = "./dataset"
IMAGES_DIR: str = os.path.join(DATASET_ROOT, "images")
ANNOTATIONS_DIR: str = os.path.join(DATASET_ROOT, "annotations")
OUTPUT_ROOT: str = "./yolo_dataset"
TRAIN_RATIO: float = 0.6
VAL_RATIO: float = 0.2
TEST_RATIO: float = 0.2
CLASS_NAMES: List[str] = ["target"]          # 通用单目标，可扩充为多类
CLASS_TO_ID: dict = {name: idx for idx, name in enumerate(CLASS_NAMES)}
# ============================================================

def parse_xml(xml_path: str) -> List[List[float]]:
    """解析 VOC XML，返回 YOLO 归一化标注列表。"""
    tree = ET.parse(xml_path)
    root = tree.getroot()
    size = root.find("size")
    img_w = float(size.find("width").text)
    img_h = float(size.find("height").text)

    objects = []
    for obj in root.iter("object"):
        cls_name = obj.find("name").text.strip()
        if cls_name not in CLASS_TO_ID:
            continue
        cls_id = CLASS_TO_ID[cls_name]
        bndbox = obj.find("bndbox")
        xmin = float(bndbox.find("xmin").text)
        ymin = float(bndbox.find("ymin").text)
        xmax = float(bndbox.find("xmax").text)
        ymax = float(bndbox.find("ymax").text)

        # 转换为 YOLO 中心坐标+宽高归一化格式
        x_center = (xmin + xmax) / 2.0 / img_w
        y_center = (ymin + ymax) / 2.0 / img_h
        width = (xmax - xmin) / img_w
        height = (ymax - ymin) / img_h
        objects.append([cls_id, x_center, y_center, width, height])
    return objects

def convert_xml_to_yolo(xml_path: str, txt_path: str) -> None:
    """将单个 XML 标注转换为 YOLO 格式 TXT 文件。"""
    objects = parse_xml(xml_path)
    with open(txt_path, "w", encoding="utf-8") as f:
        for obj in objects:
            f.write(" ".join(f"{v:.6f}" for v in obj) + "\n")

def create_dataset_structure() -> None:
    """创建 YOLO 数据集所需的目录结构。"""
    for split in ("train", "val", "test"):
        os.makedirs(os.path.join(OUTPUT_ROOT, "images", split), exist_ok=True)
        os.makedirs(os.path.join(OUTPUT_ROOT, "labels", split), exist_ok=True)

def split_and_convert() -> None:
    """主流程：匹配图像-标注对、分层划分、复制图像并转换标签。"""
    all_samples: List[Tuple[Path, Path]] = []
    for img_file in Path(IMAGES_DIR).glob("*.*"):
        if img_file.suffix.lower() not in (".jpg", ".jpeg", ".png"):
            continue
        xml_file = Path(ANNOTATIONS_DIR) / f"{img_file.stem}.xml"
        if xml_file.exists():
            all_samples.append((img_file, xml_file))

    print(f"成功匹配 {len(all_samples)} 组图像-标注对")
    if not all_samples:
        raise FileNotFoundError("未找到成对的图像与标注，请检查路径配置")

    # 分层划分：先切出 40% 临时集，再平分为验证/测试
    train_samples, temp_samples = train_test_split(
        all_samples, test_size=1 - TRAIN_RATIO, random_state=42
    )
    val_samples, test_samples = train_test_split(
        temp_samples, test_size=TEST_RATIO / (VAL_RATIO + TEST_RATIO), random_state=42
    )

    print(f"训练集: {len(train_samples)} 张, 验证集: {len(val_samples)} 张, 测试集: {len(test_samples)} 张")
    create_dataset_structure()

    for split, samples in zip(("train", "val", "test"),
                              (train_samples, val_samples, test_samples)):
        print(f"正在处理 {split} 数据集...")
        for img_path, xml_path in samples:
            # 复制图像
            shutil.copy(img_path, os.path.join(OUTPUT_ROOT, "images", split, img_path.name))
            # 生成标签
            txt_path = os.path.join(OUTPUT_ROOT, "labels", split, img_path.stem + ".txt")
            convert_xml_to_yolo(str(xml_path), txt_path)

def generate_yaml() -> None:
    """生成 Ultralytics YOLO 所需的 data.yaml 配置文件。"""
    yaml_content = (
        f"path: {os.path.abspath(OUTPUT_ROOT)}\n"
        f"train: images/train\n"
        f"val: images/val\n"
        f"test: images/test\n"
        f"nc: {len(CLASS_NAMES)}\n"
        f"names: {CLASS_NAMES}\n"
    )
    yaml_path = os.path.join(OUTPUT_ROOT, "dataset.yaml")
    with open(yaml_path, "w", encoding="utf-8") as f:
        f.write(yaml_content)
    print(f"数据集配置文件已生成：{yaml_path}")

if __name__ == "__main__":
    split_and_convert()
    generate_yaml()
    print("\n✅ 数据集转换与划分完成！")

执行转换：

python convert_and_split_dataset.py

---

三、步骤 2：YOLOv8 模型训练

训练脚本 train_yolo.py

from ultralytics import YOLO

if __name__ == "__main__":
    # 加载轻量化预训练权重，适合教学演示
    model = YOLO("yolov8n.pt")
    model.train(
        data="./yolo_dataset/dataset.yaml",
        epochs=100,
        imgsz=640,
        amp=False,          # CPU 环境建议关闭混合精度
        name="target_det"
    )

启动训练：

python train_yolo.py

训练产出位于 runs/detect/target_det，关键文件：

• results.csv / results.png：训练指标与可视化曲线
• weights/best.pt：验证集最优权重，后续推理使用

---

四、步骤 3：PySpark 分布式批量推理

核心设计思路

1. 分区复用模型 ：利用 mapPartitionsWithIndex 在每个分区仅初始化一次 YOLO 模型，消除单图加载开销；
2. 批量推理：单分区内多张图像一次性送入模型，提升 GPU/CPU 利用率；
3. 分布式汇总：各分区独立完成检测，最后在 Driver 端聚合全局目标计数；
4. 结果可视化：自动绘制检测框、标注目标数量并保存图片。

分布式推理脚本 dist_predict.py

from __future__ import annotations
import os
import cv2
from typing import Iterator, List, Tuple
from pyspark.sql import SparkSession
from ultralytics import YOLO

# ===================== 可配置参数 =====================
IMG_FOLDER: str = "./yolo_dataset/images/test"
MODEL_WEIGHT: str = "./runs/detect/target_det/weights/best.pt"
CONF_THRESH: float = 0.3          # 置信度阈值
TARGET_CLS: str = "target"        # 待检测目标类别
PARALLELISM: int = 4              # Spark 分区数
OUTPUT_DIR: str = "./outputs"
BATCH_SIZE: int = 8               # 单批推理图片数量
# =====================================================

os.makedirs(OUTPUT_DIR, exist_ok=True)

def load_model() -> YOLO:
    """每个分区仅执行一次，加载模型权重。"""
    model = YOLO(MODEL_WEIGHT)
    # 若 GPU 可用可取消注释：model.to("cuda")
    return model

def split_batch(data: list, batch_size: int) -> Iterator[list]:
    """将列表按 batch_size 切分为多个小批次。"""
    for i in range(0, len(data), batch_size):
        yield data[i:i + batch_size]

def partition_infer(part_id: int, path_iter: Iterator[str]) -> Iterator[Tuple[int, str, int]]:
    """
    分区推理回调函数。
    :param part_id: 分区编号
    :param path_iter: 该分区内的图片路径迭代器
    :return: (分区编号, 图片路径, 目标数量) 迭代器
    """
    model = load_model()
    all_paths: List[str] = []
    all_imgs: List[cv2.Mat] = []

    # 读取当前分区全部图像
    for p in path_iter:
        all_paths.append(p)
        img = cv2.imread(p)
        if img is not None:
            all_imgs.append(img)
        else:
            # 读图失败时保持长度一致
            all_imgs.append(None)

    # 分批推理
    path_batches = list(split_batch(all_paths, BATCH_SIZE))
    img_batches = list(split_batch(all_imgs, BATCH_SIZE))

    results: List[Tuple[int, str, int]] = []
    for paths, imgs in zip(path_batches, img_batches):
        valid_indices = [i for i, img in enumerate(imgs) if img is not None]
        if not valid_indices:
            continue
        # 只对有效图像进行推理
        valid_imgs = [imgs[i] for i in valid_indices]
        preds = model(valid_imgs, conf=CONF_THRESH)

        for i, pred in zip(valid_indices, preds):
            img_src = imgs[i].copy()
            target_count = 0
            if pred.boxes is not None:
                boxes = pred.boxes.xyxy.cpu().numpy()
                classes = pred.boxes.cls.cpu().numpy()
                names = pred.names
                for box, cls_id in zip(boxes, classes):
                    if names[int(cls_id)] == TARGET_CLS:
                        target_count += 1
                        x1, y1, x2, y2 = map(int, box)
                        cv2.rectangle(img_src, (x1, y1), (x2, y2), (0, 255, 0), 2)

            # 标注目标总数
            cv2.putText(img_src, f"Total: {target_count}", (20, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

            # 保存结果图像
            img_name = os.path.basename(paths[i])
            save_path = os.path.join(OUTPUT_DIR, img_name)
            cv2.imwrite(save_path, img_src)
            results.append((part_id, paths[i], target_count))

    return iter(results)


if __name__ == "__main__":
    spark = SparkSession.builder \
        .appName("Yolo_Dist_Detect") \
        .master("local[*]") \
        .getOrCreate()
    sc = spark.sparkContext
    sc.setLogLevel("WARN")

    # 收集所有测试图像绝对路径
    img_paths = []
    for fname in os.listdir(IMG_FOLDER):
        if fname.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
            img_paths.append(os.path.abspath(os.path.join(IMG_FOLDER, fname)))

    # 构建 RDD 并执行分布式推理
    rdd = sc.parallelize(img_paths, numSlices=PARALLELISM)
    result_rdd = rdd.mapPartitionsWithIndex(partition_infer)
    all_results = result_rdd.collect()

    # 输出统计信息
    print(f"===== 分布式图像检测 | 分区数: {PARALLELISM} | 单批大小: {BATCH_SIZE} =====")
    total_targets = 0
    for pid, path, num in all_results:
        fname = os.path.basename(path)
        print(f"【分区:{pid:2d}】图像:{fname:22s} 目标数量:{num}")
        total_targets += num
    print(f"\n全部图像累计检测目标总数：{total_targets}")

    spark.stop()

运行推理：

python dist_predict.py

控制台输出示例：

===== 分布式图像检测 | 分区数:4 | 单批大小:8 =====
【分区: 0】图像:img_01947.jpg        目标数量:69
【分区: 0】图像:img_00717.jpg        目标数量:5
【分区: 1】图像:img_00041.jpg        目标数量:10
...
全部图像累计检测目标总数：1256

推理结果可视化图片保存在 outputs 目录下。

---

五、核心知识点梳理

1. 大数据分布式编程要点

• mapPartitionsWithIndex：按分区批量操作，可获取分区 ID，适合模型、数据库连接等重型资源的初始化；
• RDD 并行化 ：通过 sc.parallelize 的 numSlices 控制并行度；
• 惰性执行 ：转换算子仅记录血缘，collect() 触发实际计算；
• 分布式聚合：各分区独立计算，最终在 Driver 端完成全局统计。

2. YOLO 目标检测关键环节

• VOC 坐标到 YOLO 归一化格式的转换原理；
• 数据集分层划分，防止数据分布偏移；
• 批量推理以提升硬件利用率；
• 置信度过滤、类别筛选、检测框绘制方法。

3. 大数据 + AI 融合避坑指南

• ❌ 切忌在 map 内加载模型，否则每张图片都会重新初始化，开销巨大；
• ✅ 应在分区入口函数加载模型，每个分区仅初始化一次；
• ✅ 分布式环境下务必使用绝对路径读取文件；
• ✅ 建议将 Spark 日志级别设为 WARN，避免海量 INFO 日志干扰结果查看；
• ✅ 根据实际内存/显存灵活调整 BATCH_SIZE。

---

六、拓展优化方向

1. 对接 HDFS，支持从分布式存储读取海量图像；
2. 引入 Kafka 构建流式实时目标检测管道；
3. 将检测结果写入关系型数据库，搭建可视化看板；
4. 设计单机与分布式推理的性能对比实验，量化加速比；
5. 扩展多类别检测，仅需修改 CLASS_NAMES 配置即可通用。

---

七、总结

本项目提供了一个无业务场景锁定的通用模板，完整串联了"数据预处理 → 模型训练 → 分布式推理"全链路。重点演示了分区资源复用 与分布式结果聚合两大核心思想，适用于大数据与 AI 交叉方向的教学、练习与快速原型开发。