Z-Image-Turbo工业检测:YOLOv5集成方案
想象一下,工厂流水线上,质检员正盯着屏幕,一张张检查产品外观。划痕、污渍、尺寸偏差……眼睛看花了,效率还上不去,漏检率也居高不下。这种场景在很多制造企业里每天都在上演。
现在,情况正在改变。我们尝试把阿里开源的Z-Image-Turbo图像生成模型,和经典的YOLOv5目标检测算法结合起来,搭建一套智能化的工业产品质量自动检测系统。听起来有点技术含量,但实际用起来,你会发现它比想象中要简单得多。
1. 为什么要把Z-Image-Turbo和YOLOv5放一起?
先说说这两个技术各自是干什么的。
YOLOv5你可能听说过,它是目前工业界用得最多的目标检测算法之一。简单来说,就是让计算机“看”图片,然后告诉你图片里有什么东西,在什么位置。比如一张电路板的照片,YOLOv5能识别出上面的电容、电阻、芯片,还能标出它们的具体位置。
Z-Image-Turbo是阿里最近开源的一个图像生成模型。它的特点是“快”和“轻”——只需要8步就能生成一张高质量的图片,而且对硬件要求不高,普通消费级显卡就能跑起来。
你可能会问:一个生成图片的模型,和一个检测图片的模型,放一起能干什么?
这里面的逻辑其实很巧妙。在工业检测场景里,我们经常遇到一个问题:缺陷样本太少了。一个正常运转的生产线,99%的产品都是合格的,只有1%甚至更少的产品有缺陷。你想训练一个能准确识别缺陷的AI模型,但手头只有几十张、几百张缺陷图片,根本不够用。
这时候Z-Image-Turbo的价值就体现出来了。我们可以用它来生成各种缺陷样本——不同位置、不同大小、不同形状的划痕、污渍、变形等等。有了这些生成的样本,我们就能大大扩充YOLOv5的训练数据集,让它学得更准、更稳。
1.1 传统方法的痛点
以前遇到样本不足的问题,大家通常用几种方法:
- 数据增强:把现有的图片旋转、缩放、加噪声,但本质上还是那些图片,多样性有限
- 人工标注:找人来一张张标注,成本高、效率低,还容易出错
- 迁移学习:用其他场景训练好的模型来微调,但不同产品差异大,效果不一定好
这些方法要么效果有限,要么成本太高。而用Z-Image-Turbo生成缺陷样本,相当于请了一个“虚拟质检员”,24小时不间断地“创造”各种缺陷情况,而且完全免费。
1.2 我们的解决方案思路
我们的方案分三步走:
- 用Z-Image-Turbo生成缺陷样本:根据产品特点,描述各种可能的缺陷情况,批量生成带缺陷的产品图片
- 用YOLOv5训练检测模型:把生成的缺陷图片和真实的正品图片混合,训练一个专门识别缺陷的模型
- 部署到生产线:把训练好的模型部署到实际的生产环境中,实现自动化检测
整个过程就像教一个新手质检员:先给他看大量“教材”(生成的缺陷样本),让他熟悉各种情况,然后再去实际岗位工作。
2. 环境准备与快速部署
2.1 硬件要求
这套方案对硬件的要求比较友好:
- 显卡:NVIDIA GPU,显存8GB以上(RTX 3060/3070/4060等都可以)
- 内存:16GB以上
- 存储:至少50GB可用空间(主要存放模型文件和数据集)
如果你没有独立显卡,用CPU也能跑,只是速度会慢一些。对于测试和小规模应用来说,完全够用。
2.2 软件环境搭建
我们先来搭建基础环境。这里用Python 3.8+和PyTorch作为主要框架。
# 创建虚拟环境(可选但推荐) python -m venv zimage_yolo_env source zimage_yolo_env/bin/activate # Linux/Mac # 或者 zimage_yolo_env\Scripts\activate # Windows # 安装PyTorch(根据你的CUDA版本选择) pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 安装其他依赖 pip install diffusers transformers accelerate pillow opencv-python pip install ultralytics # YOLOv5 pip install dashscope # 阿里云API(如果用API方式)2.3 Z-Image-Turbo模型下载
Z-Image-Turbo有几种不同的版本,我们选择最适合工业场景的版本:
import os from diffusers import DiffusionPipeline import torch # 模型下载和加载 def load_z_image_turbo(): """ 加载Z-Image-Turbo模型 返回:生成管道 """ # 使用阿里云ModelScope的模型 model_id = "Tongyi-MAI/Z-Image-Turbo" # 创建管道 pipe = DiffusionPipeline.from_pretrained( model_id, torch_dtype=torch.bfloat16, # 使用bfloat16减少显存占用 safety_checker=None, # 工业场景不需要安全检测 requires_safety_checker=False ) # 启用CPU卸载(如果显存不足) pipe.enable_model_cpu_offload() # 启用Flash Attention加速(如果显卡支持) try: pipe.transformer.set_attention_backend("flash") print("Flash Attention已启用") except: print("Flash Attention不可用,使用默认注意力机制") return pipe # 测试模型加载 print("正在加载Z-Image-Turbo模型...") z_image_pipe = load_z_image_turbo() print("模型加载完成!")如果你显存比较紧张(比如只有8GB),可以用量化版本:
# 使用量化版本(显存占用更低) def load_z_image_turbo_quantized(): from diffusers import DPMSolverMultistepScheduler pipe = DiffusionPipeline.from_pretrained( "Tongyi-MAI/Z-Image-Turbo", torch_dtype=torch.float16, # 使用float16进一步减少显存 variant="fp16", # 使用fp16变体 ) # 使用更快的调度器 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config) # 启用CPU卸载 pipe.enable_model_cpu_offload() return pipe3. 生成工业缺陷样本
这是整个方案的核心部分。我们要用Z-Image-Turbo生成各种产品缺陷的图片。
3.1 理解工业缺陷类型
不同的产品有不同的缺陷类型,但大体上可以分为几类:
- 表面缺陷:划痕、污渍、色差、气泡、凹陷
- 尺寸缺陷:过大、过小、变形、不对称
- 装配缺陷:零件缺失、错位、松动
- 功能缺陷:这个需要实际测试,图片看不出来
我们先以最常见的电子产品外壳为例,看看怎么生成缺陷样本。
3.2 基础缺陷生成代码
import torch from PIL import Image import numpy as np import os def generate_defect_samples(product_type, defect_type, num_samples=10): """ 生成指定产品和缺陷类型的样本 参数: product_type: 产品类型,如"手机外壳"、"电路板" defect_type: 缺陷类型,如"划痕"、"污渍" num_samples: 生成样本数量 """ # 根据产品和缺陷类型构造提示词 prompts = construct_prompts(product_type, defect_type, num_samples) generated_images = [] for i, prompt in enumerate(prompts): print(f"生成第 {i+1}/{num_samples} 张图片...") # 生成图片 image = z_image_pipe( prompt=prompt, negative_prompt="完美无瑕,无缺陷,高质量,专业摄影", # 反向提示词,避免生成完美图片 guidance_scale=0.0, # Turbo模型必须设为0 num_inference_steps=8, # Turbo模型只需要8步 height=1024, width=1024, generator=torch.Generator(device="cuda").manual_seed(i) # 固定种子保证可重复性 ).images[0] # 保存图片 save_path = f"defect_samples/{product_type}/{defect_type}/sample_{i}.png" os.makedirs(os.path.dirname(save_path), exist_ok=True) image.save(save_path) generated_images.append((save_path, prompt)) return generated_images def construct_prompts(product_type, defect_type, num_samples): """ 构造生成提示词 """ prompts = [] # 基础描述 base_descriptions = { "手机外壳": "一个智能手机的金属外壳,表面光滑,", "电路板": "一块绿色的印刷电路板,上面有电子元件,", "塑料零件": "一个注塑成型的塑料零件,", "玻璃面板": "一块手机玻璃面板," } # 缺陷描述 defect_descriptions = { "划痕": [ "表面有一条明显的划痕,长度约2-3厘米", "有多条交叉的细微划痕", "深度划痕,能看到底层材料", "弧形划痕,像是被尖锐物体刮过" ], "污渍": [ "表面有油污污渍,反光不均匀", "有指纹印迹,清晰可见", "有水渍痕迹,形成环状", "有灰尘积聚,特别是在边缘处" ], "气泡": [ "表面有微小气泡,直径约1-2毫米", "多个气泡聚集在一起", "大气泡,直径超过5毫米", "气泡破裂后的痕迹" ], "色差": [ "颜色不均匀,局部偏深", "与标准色有明显差异", "表面有斑驳的色块", "边缘处颜色变浅" ] } # 视角和光线 views = ["正面视角", "侧面视角", "45度角", "俯视角度"] lighting = ["均匀光照", "侧光照明", "顶光照明", "漫反射光"] # 生成多样化的提示词 for i in range(num_samples): base = base_descriptions.get(product_type, "一个工业产品,") defect = defect_descriptions.get(defect_type, ["有缺陷"])[i % len(defect_descriptions[defect_type])] view = views[i % len(views)] light = lighting[i % len(lighting)] prompt = f"{base}{defect},{view},{light},工业摄影风格,高清细节,白色背景" prompts.append(prompt) return prompts # 示例:生成手机外壳划痕样本 print("开始生成缺陷样本...") samples = generate_defect_samples("手机外壳", "划痕", num_samples=5) print(f"生成了 {len(samples)} 个样本")3.3 批量生成与数据增强
实际应用中,我们需要生成大量样本,并且要做一些后处理:
import cv2 from PIL import Image, ImageFilter, ImageEnhance import random def augment_defect_image(image_path, augmentations=None): """ 对生成的缺陷图片进行数据增强 参数: image_path: 图片路径 augmentations: 增强操作列表,如["rotate", "blur", "noise"] """ if augmentations is None: augmentations = ["rotate", "brightness", "contrast", "blur"] image = Image.open(image_path) # 随机选择几种增强方式 selected_augs = random.sample(augmentations, k=min(3, len(augmentations))) for aug in selected_augs: if aug == "rotate" and random.random() > 0.5: # 随机旋转 angle = random.uniform(-10, 10) image = image.rotate(angle, expand=True, fillcolor=(255, 255, 255)) elif aug == "brightness": # 调整亮度 factor = random.uniform(0.8, 1.2) enhancer = ImageEnhance.Brightness(image) image = enhancer.enhance(factor) elif aug == "contrast": # 调整对比度 factor = random.uniform(0.8, 1.2) enhancer = ImageEnhance.Contrast(image) image = enhancer.enhance(factor) elif aug == "blur" and random.random() > 0.7: # 轻微模糊(模拟对焦不准) image = image.filter(ImageFilter.GaussianBlur(radius=random.uniform(0.5, 1.5))) elif aug == "noise": # 添加噪声 img_array = np.array(image) noise = np.random.normal(0, random.uniform(1, 3), img_array.shape) img_array = np.clip(img_array + noise, 0, 255).astype(np.uint8) image = Image.fromarray(img_array) return image def batch_generate_with_augmentation(product_type, defect_types, samples_per_type=20, augment_per_image=3): """ 批量生成并增强缺陷样本 """ all_samples = [] for defect_type in defect_types: print(f"生成 {product_type} 的 {defect_type} 缺陷样本...") # 生成基础样本 base_samples = generate_defect_samples(product_type, defect_type, samples_per_type) # 对每个样本进行增强 for img_path, prompt in base_samples: all_samples.append((img_path, prompt, defect_type, "original")) # 生成增强版本 for aug_idx in range(augment_per_image): augmented_img = augment_defect_image(img_path) # 保存增强版本 aug_path = img_path.replace(".png", f"_aug{aug_idx}.png") augmented_img.save(aug_path) all_samples.append((aug_path, prompt + f" [增强{aug_idx}]", defect_type, f"aug{aug_idx}")) print(f"总共生成了 {len(all_samples)} 个样本(含增强)") return all_samples # 批量生成多种缺陷 defect_types = ["划痕", "污渍", "气泡", "色差"] all_samples = batch_generate_with_augmentation("手机外壳", defect_types, samples_per_type=10, augment_per_image=2)4. YOLOv5模型训练
有了足够的缺陷样本,我们就可以开始训练YOLOv5检测模型了。
4.1 数据准备与标注
YOLOv5需要标注数据,也就是告诉模型缺陷在图片的什么位置。我们可以用半自动的方式来完成:
import json from pathlib import Path import yaml def create_yolo_dataset(samples, output_dir="yolo_dataset"): """ 创建YOLOv5格式的数据集 """ # 创建目录结构 dataset_dir = Path(output_dir) (dataset_dir / "images" / "train").mkdir(parents=True, exist_ok=True) (dataset_dir / "images" / "val").mkdir(parents=True, exist_ok=True) (dataset_dir / "labels" / "train").mkdir(parents=True, exist_ok=True) (dataset_dir / "labels" / "val").mkdir(parents=True, exist_ok=True) # 缺陷类别映射 defect_classes = { "划痕": 0, "污渍": 1, "气泡": 2, "色差": 3 } # 分割训练集和验证集 train_ratio = 0.8 split_idx = int(len(samples) * train_ratio) train_samples = samples[:split_idx] val_samples = samples[split_idx:] def process_samples(sample_list, split_name): for idx, (img_path, prompt, defect_type, aug_type) in enumerate(sample_list): # 复制图片 img = Image.open(img_path) new_img_path = dataset_dir / "images" / split_name / f"{split_name}_{idx}.jpg" img.save(new_img_path) # 创建标注文件(这里简化处理,实际需要标注工具) # 假设缺陷位于图片中心区域(实际应用中需要用标注工具) label_path = dataset_dir / "labels" / split_name / f"{split_name}_{idx}.txt" # 生成模拟的标注(中心区域) with open(label_path, 'w') as f: class_id = defect_classes.get(defect_type, 0) # 假设缺陷在图片中心,占图片的30% x_center = 0.5 y_center = 0.5 width = 0.3 height = 0.3 f.write(f"{class_id} {x_center} {y_center} {width} {height}\n") # 处理训练集和验证集 print("处理训练集...") process_samples(train_samples, "train") print("处理验证集...") process_samples(val_samples, "val") # 创建数据集配置文件 create_dataset_yaml(dataset_dir, defect_classes) return dataset_dir def create_dataset_yaml(dataset_dir, defect_classes): """ 创建YOLOv5数据集配置文件 """ config = { "path": str(dataset_dir.absolute()), "train": "images/train", "val": "images/val", "names": {str(k): v for v, k in enumerate(defect_classes.keys())} } config_path = dataset_dir / "dataset.yaml" with open(config_path, 'w') as f: yaml.dump(config, f, default_flow_style=False) print(f"数据集配置文件已创建: {config_path}") return config_path # 创建YOLOv5数据集 print("创建YOLOv5数据集...") dataset_dir = create_yolo_dataset(all_samples) print(f"数据集已创建到: {dataset_dir}")4.2 YOLOv5模型训练
现在我们可以开始训练了。YOLOv5提供了非常简单的训练接口:
from ultralytics import YOLO import torch def train_yolov5_model(dataset_config, model_size="s", epochs=50): """ 训练YOLOv5模型 参数: dataset_config: 数据集配置文件路径 model_size: 模型大小,可选 "n", "s", "m", "l", "x" epochs: 训练轮数 """ # 加载预训练模型 print(f"加载YOLOv5{model_size}预训练模型...") model = YOLO(f"yolov5{model_size}.pt") # 训练配置 train_args = { "data": dataset_config, "epochs": epochs, "imgsz": 640, "batch": 16, "workers": 4, "device": "0" if torch.cuda.is_available() else "cpu", "project": "defect_detection", "name": f"zimage_yolov5_{model_size}", "exist_ok": True, "pretrained": True, "optimizer": "AdamW", "lr0": 0.001, "cos_lr": True, # 使用余弦学习率调度 "label_smoothing": 0.1, "patience": 20, # 早停耐心值 "save_period": 10, # 每10轮保存一次 "seed": 42 } # 开始训练 print("开始训练YOLOv5模型...") results = model.train(**train_args) print("训练完成!") return model, results # 开始训练 print("开始训练缺陷检测模型...") model, results = train_yolov5_model(str(dataset_dir / "dataset.yaml"), model_size="s", epochs=50)4.3 训练过程监控
训练过程中,我们可以监控一些关键指标:
import matplotlib.pyplot as plt def plot_training_results(results): """ 绘制训练结果图表 """ # 这里results对象包含了训练过程中的各种指标 # 实际使用时需要根据YOLOv5的输出格式调整 fig, axes = plt.subplots(2, 2, figsize=(12, 8)) # 损失曲线 axes[0, 0].plot(results['train_loss'], label='训练损失') axes[0, 0].plot(results['val_loss'], label='验证损失') axes[0, 0].set_title('损失曲线') axes[0, 0].set_xlabel('轮次') axes[0, 0].set_ylabel('损失') axes[0, 0].legend() axes[0, 0].grid(True) # 准确率曲线 axes[0, 1].plot(results['precision'], label='精确率') axes[0, 1].plot(results['recall'], label='召回率') axes[0, 1].set_title('精确率/召回率') axes[0, 1].set_xlabel('轮次') axes[0, 1].set_ylabel('百分比') axes[0, 1].legend() axes[0, 1].grid(True) # mAP曲线 axes[1, 0].plot(results['map50'], label='mAP@0.5') axes[1, 0].plot(results['map50_95'], label='mAP@0.5:0.95') axes[1, 0].set_title('mAP指标') axes[1, 0].set_xlabel('轮次') axes[1, 0].set_ylabel('mAP') axes[1, 0].legend() axes[1, 0].grid(True) # 学习率曲线 axes[1, 1].plot(results['learning_rate'], label='学习率') axes[1, 1].set_title('学习率变化') axes[1, 1].set_xlabel('轮次') axes[1, 1].set_ylabel('学习率') axes[1, 1].legend() axes[1, 1].grid(True) plt.tight_layout() plt.savefig('training_results.png', dpi=150, bbox_inches='tight') plt.show() # 注意:实际results对象的结构可能需要调整 # plot_training_results(results)5. 实际应用与部署
模型训练好后,我们要把它部署到实际的生产环境中。
5.1 模型验证与测试
先测试一下模型的效果:
def test_defect_detection(model_path, test_images_dir): """ 测试缺陷检测模型 """ # 加载训练好的模型 model = YOLO(model_path) # 测试图片目录 test_images = list(Path(test_images_dir).glob("*.jpg")) + list(Path(test_images_dir).glob("*.png")) results = [] for img_path in test_images[:10]: # 测试前10张 print(f"检测: {img_path.name}") # 运行检测 detection_results = model(img_path, conf=0.25, iou=0.45) # 解析结果 for result in detection_results: boxes = result.boxes if boxes is not None: for box in boxes: cls = int(box.cls[0]) conf = float(box.conf[0]) xyxy = box.xyxy[0].tolist() results.append({ "image": img_path.name, "class": cls, "confidence": conf, "bbox": xyxy }) print(f" 检测到缺陷: 类别={cls}, 置信度={conf:.3f}, 位置={xyxy}") else: print(f" 未检测到缺陷") return results # 测试模型 print("测试缺陷检测模型...") test_results = test_defect_detection("defect_detection/zimage_yolov5_s/weights/best.pt", "test_images")5.2 部署到生产线
在实际生产线上,我们需要一个完整的检测系统:
import cv2 import time from datetime import datetime import json class DefectDetectionSystem: """ 缺陷检测系统 """ def __init__(self, model_path, camera_index=0): """ 初始化检测系统 参数: model_path: YOLOv5模型路径 camera_index: 摄像头索引 """ # 加载模型 self.model = YOLO(model_path) # 初始化摄像头 self.cap = cv2.VideoCapture(camera_index) # 检测阈值 self.conf_threshold = 0.3 self.iou_threshold = 0.4 # 统计信息 self.total_inspected = 0 self.defect_count = 0 self.start_time = time.time() # 缺陷类别名称 self.class_names = { 0: "划痕", 1: "污渍", 2: "气泡", 3: "色差" } def process_frame(self, frame): """ 处理一帧图像 """ # 运行检测 results = self.model(frame, conf=self.conf_threshold, iou=self.iou_threshold) detected_defects = [] for result in results: boxes = result.boxes if boxes is not None: for box in boxes: cls = int(box.cls[0]) conf = float(box.conf[0]) xyxy = box.xyxy[0].tolist() defect_info = { "type": self.class_names.get(cls, "未知"), "confidence": conf, "bbox": xyxy, "timestamp": datetime.now().isoformat() } detected_defects.append(defect_info) # 在图像上绘制检测框 x1, y1, x2, y2 = map(int, xyxy) cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2) label = f"{self.class_names.get(cls, '未知')}: {conf:.2f}" cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) return frame, detected_defects def run_realtime_detection(self): """ 运行实时检测 """ print("启动实时缺陷检测系统...") print("按 'q' 键退出") while True: # 读取帧 ret, frame = self.cap.read() if not ret: print("无法读取摄像头") break # 处理帧 processed_frame, defects = self.process_frame(frame) # 更新统计 self.total_inspected += 1 if defects: self.defect_count += 1 # 显示统计信息 stats_text = f"检测数: {self.total_inspected} | 缺陷数: {self.defect_count}" cv2.putText(processed_frame, stats_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) # 显示结果 cv2.imshow('Defect Detection System', processed_frame) # 保存检测记录 if defects: self.save_detection_record(defects, frame) # 检查退出 if cv2.waitKey(1) & 0xFF == ord('q'): break # 清理 self.cap.release() cv2.destroyAllWindows() # 打印最终统计 self.print_statistics() def save_detection_record(self, defects, frame): """ 保存检测记录 """ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") # 保存图像 img_path = f"detection_records/{timestamp}.jpg" os.makedirs(os.path.dirname(img_path), exist_ok=True) cv2.imwrite(img_path, frame) # 保存检测结果 record = { "timestamp": timestamp, "image_path": img_path, "defects": defects, "total_defects": len(defects) } record_path = f"detection_records/{timestamp}.json" with open(record_path, 'w') as f: json.dump(record, f, indent=2) def print_statistics(self): """ 打印统计信息 """ elapsed_time = time.time() - self.start_time print("\n" + "="*50) print("检测系统统计") print("="*50) print(f"运行时间: {elapsed_time:.1f} 秒") print(f"检测总数: {self.total_inspected}") print(f"缺陷数量: {self.defect_count}") if self.total_inspected > 0: defect_rate = (self.defect_count / self.total_inspected) * 100 print(f"缺陷率: {defect_rate:.2f}%") print(f"检测速度: {self.total_inspected/elapsed_time:.1f} 帧/秒") print("="*50) # 运行检测系统(需要摄像头) # detection_system = DefectDetectionSystem("defect_detection/zimage_yolov5_s/weights/best.pt") # detection_system.run_realtime_detection()5.3 批量处理与报告生成
对于离线批量检测,我们可以这样处理:
import pandas as pd from tqdm import tqdm def batch_process_images(image_dir, model_path, output_dir="batch_results"): """ 批量处理图片 """ # 创建输出目录 os.makedirs(output_dir, exist_ok=True) # 加载模型 model = YOLO(model_path) # 获取所有图片 image_files = [] for ext in ["*.jpg", "*.jpeg", "*.png", "*.bmp"]: image_files.extend(Path(image_dir).glob(ext)) all_results = [] print(f"开始批量处理 {len(image_files)} 张图片...") for img_path in tqdm(image_files): # 读取图片 img = cv2.imread(str(img_path)) if img is None: continue # 运行检测 results = model(img, conf=0.25) image_results = { "filename": img_path.name, "defect_count": 0, "defect_types": [], "max_confidence": 0 } for result in results: boxes = result.boxes if boxes is not None: defects_in_image = [] for box in boxes: cls = int(box.cls[0]) conf = float(box.conf[0]) defect_type = ["划痕", "污渍", "气泡", "色差"][cls] if cls < 4 else f"未知{cls}" defects_in_image.append({ "type": defect_type, "confidence": conf }) image_results["defect_count"] = len(defects_in_image) image_results["defect_types"] = [d["type"] for d in defects_in_image] if defects_in_image: image_results["max_confidence"] = max(d["confidence"] for d in defects_in_image) # 保存带标注的图片 annotated_img = result.plot() output_path = Path(output_dir) / f"annotated_{img_path.name}" cv2.imwrite(str(output_path), annotated_img) all_results.append(image_results) # 生成报告 generate_report(all_results, output_dir) return all_results def generate_report(results, output_dir): """ 生成检测报告 """ df = pd.DataFrame(results) # 基本统计 total_images = len(df) defective_images = df[df["defect_count"] > 0].shape[0] defect_rate = (defective_images / total_images * 100) if total_images > 0 else 0 # 缺陷类型分布 defect_types = [] for types in df["defect_types"]: defect_types.extend(types) type_counts = pd.Series(defect_types).value_counts() # 生成报告文本 report = f""" 缺陷检测报告 生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} {'='*50} 总体统计: 总图片数: {total_images} 缺陷图片数: {defective_images} 缺陷率: {defect_rate:.2f}% 缺陷类型分布: {type_counts.to_string()} 详细结果: {df.to_string(index=False)} """ # 保存报告 report_path = Path(output_dir) / "detection_report.txt" with open(report_path, 'w', encoding='utf-8') as f: f.write(report) # 保存CSV csv_path = Path(output_dir) / "detection_results.csv" df.to_csv(csv_path, index=False, encoding='utf-8-sig') print(f"报告已生成: {report_path}") print(f"CSV数据已保存: {csv_path}") return report # 批量处理示例 # batch_results = batch_process_images("production_images", "defect_detection/zimage_yolov5_s/weights/best.pt")6. 系统优化与扩展
6.1 性能优化建议
实际部署时,可以考虑以下优化:
def optimize_for_production(model_path): """ 为生产环境优化模型 """ # 加载模型 model = YOLO(model_path) # 1. 模型量化(减少内存和计算量) quantized_model = model.quantize() # 2. 转换为ONNX格式(提高跨平台兼容性) model.export(format="onnx", imgsz=640, simplify=True) # 3. 使用TensorRT加速(如果使用NVIDIA GPU) model.export(format="engine", imgsz=640) print("模型优化完成") return quantized_model # 4. 使用多线程处理 from concurrent.futures import ThreadPoolExecutor import threading class ParallelDetectionSystem: """ 并行检测系统 """ def __init__(self, model_path, num_workers=4): self.model = YOLO(model_path) self.executor = ThreadPoolExecutor(max_workers=num_workers) self.lock = threading.Lock() self.results = [] def detect_single(self, image_path): """单张图片检测""" results = self.model(image_path) return results def detect_batch_parallel(self, image_paths): """并行批量检测""" futures = [] for img_path in image_paths: future = self.executor.submit(self.detect_single, img_path) futures.append(future) # 收集结果 all_results = [] for future in futures: all_results.append(future.result()) return all_results6.2 扩展到其他工业场景
这套方案可以很容易地扩展到其他工业场景:
- 纺织品检测:生成各种布料缺陷(破洞、污渍、色差)
- 食品检测:生成食品外观缺陷(形状异常、颜色不均)
- 药品检测:生成药片缺陷(破损、尺寸不符)
- 汽车零件检测:生成零件表面缺陷
只需要调整生成样本时的提示词和训练数据即可。
7. 总结
把Z-Image-Turbo和YOLOv5结合起来做工业检测,实际用下来效果比预想的要好。Z-Image-Turbo生成缺陷样本的能力确实强,特别是对中文提示词的理解很到位,能生成符合我们需求的缺陷图片。YOLOv5作为检测框架,成熟稳定,训练和部署都很方便。
这套方案最大的优势是解决了工业检测中“缺陷样本少”的痛点。以前要收集几百张缺陷图片可能需要几个月时间,现在用Z-Image-Turbo,一两天就能生成足够多的样本。而且生成样本的多样性很好,各种位置、各种大小的缺陷都能覆盖到。
部署方面,YOLOv5的生态很完善,从训练到部署的整套工具链都很成熟。我们可以在本地训练好模型,然后轻松部署到生产线上的工控机、边缘设备甚至云端服务器。
当然,实际应用中还会遇到一些具体问题,比如生成样本和真实样本的分布差异、不同产品需要不同的提示词策略等。但这些都可以通过迭代优化来解决。整体来说,这套方案为工业质检提供了一个成本低、效果好的AI解决方案。
如果你也在做工业检测相关的项目,不妨试试这个组合。从生成样本到训练模型,再到部署应用,整个流程跑通后,你会对AI在工业领域的应用有更直观的感受。而且随着Z-Image-Turbo这类模型的不断优化,未来的效果还会更好。
获取更多AI镜像
想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。