HSG金属表面缺陷检测
- 1. 项目背景
- 1.1 项目简述
- 1.2 项目目标
- 2. 解决方案
- 3. 数据集
- 3.1 收集各种缺陷的图片
- 3.2 利用有限图片创造更多可能
- 3.3 分割图像
- 3.4 打标签
- 4. 部分代码
- 4.1 数据集划分
- 4.2图像分割
- 4.3 训练模型
- 4.4 预测
- 5. 预测结果
1. 项目背景
1.1 项目简述
- iPad HSG 的材质是带有磨砂质感的金属框，在生产或搬运过程中，会产生一些缺陷。比如：长度不等的黑线，面积不定的亮印，亮度不同的吐酸酸蚀，以及HSG在搬运过程中经常出现的刮伤或磕伤。
- 常见的4种缺陷类型：
-
划伤
-
吐酸
-
黑线
-
亮印
-
1.2 项目目标
- 使用多组相机对HSG进行拍照取像，通过对图像的分析，完成以上四种缺陷的检测。
- 技术要求：明显缺陷要求 100% 检出，不明显缺陷检出率超过 95%。
2. 解决方案
- 该项目遇到的问题属于检测类问题，最终选择使用 YOLOv8 Detect 模块完成。
3. 数据集
3.1 收集各种缺陷的图片
- 自动化领域的缺陷样品收集起来并不容易，因为带有缺陷的样品往往当天或者第二天就会被返工或做报废处理，工程师到达现场时只能拿到当天的一小部分缺陷样品。在打样阶段产品本来就少，这给我们的工作带来不少麻烦。
3.2 利用有限图片创造更多可能
- 能拿到的样品少，就只能在实验室制造更多可能了。
- 改变拍摄的位置和角度，光源亮度，拍摄更多的图片。
- 也可以通过图像处理（旋转、翻转等）生成更多的图片。
3.3 分割图像
- 项目使用的是 1200W 彩色相机，缺陷相较于整张图像来说比较小，所以采取了图像分割的方法。将原有的大图像分割为小图，再进行标注。
- 这里采用 5 × 4 分割。
3.4 打标签
- 标注工具使用的是 lableImg。
- calsses：
- acid
- bright
- black
- scratch
4. 部分代码
4.1 数据集划分
import cv2
import matplotlib.pyplot as plt
import numpy as np
import random
from tqdm import tqdm
import shutil
import osdef CollateDataset(image_dir,label_dir):"""功能：数据集划分（训练集、验证集、测试集）:param image_dir: 图片路径:param label_dir: 标签路径:return:"""# 创建一个空列表来存储有效图片的路径valid_images = []# 创建一个空列表来存储有效 label 的路径valid_labels = []# 遍历 images 文件夹下的所有图片for image_name in os.listdir(image_dir):# 获取图片的完整路径image_path = os.path.join(image_dir, image_name)# 获取图片文件的扩展名ext = os.path.splitext(image_name)[-1]# 根据扩展名替换成对应的 label 文件名label_name = image_name.replace(ext, ".txt")# 获取对应 label 的完整路径label_path = os.path.join(label_dir, label_name)# 判断 label 是否存在if not os.path.exists(label_path):# # 删除图片# os.remove(image_path)print("there is no:", label_path)else:# 将图片路径添加到列表中valid_images.append(image_path)# 将 label 路径添加到列表中valid_labels.append(label_path)# 遍历每个有效图片路径for i in tqdm(range(len(valid_images))):image_path = valid_images[i]label_path = valid_labels[i]# 随机生成一个概率r = random.random()# 判断图片应该移动到哪个文件夹# train：valid：test = 8:2:0if r < 0.0:# 移动到 test 文件夹destination = "./dataset/test"elif r < 0.2:# 移动到 valid 文件夹destination = "./dataset/valid"else:# 移动到 train 文件夹destination = "./dataset/train"# 创建目标文件夹中 images 和 labels 子文件夹os.makedirs(os.path.join(destination, "images"), exist_ok=True)os.makedirs(os.path.join(destination, "labels"), exist_ok=True)# 生成目标文件夹中图片的新路径image_destination_path = os.path.join(destination, "images", os.path.basename(image_path))# 移动图片到目标文件夹shutil.copy(image_path, image_destination_path)# 生成目标文件夹中 label 的新路径label_destination_path = os.path.join(destination, "labels", os.path.basename(label_path))# 移动 label 到目标文件夹shutil.copy(label_path, label_destination_path)if __name__ == '__main__':CollateDataset("./images","./labels")
4.2图像分割
import cv2
import os
import pathlibsource_path = "D:/images"
target_path = "D:/images_split"def mkdir(path):path = path.strip()path = path.rstrip("\\")isExists = os.path.exists(path)if not isExists:os.makedirs(path)print(path + ' 创建成功')return Trueelse:print(path + ' 目录已存在')return Falsedef split_img2(img_file):img = cv2.imread(img_file)img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)img_h = img.shape[0] # 高度img_w = img.shape[1] # 宽度# h1_half = int(img_h / 2)# w1_half = int(img_w / 2)h1_half = img_h // 2w1_half = img_w // 2img_name = os.path.basename(img_file)for i in range(4):img1 = img[int(i / 2) * h1_half: h1_half * (int(i / 2) + 1), int(i % 2) * w1_half: (int(i % 2) + 1) * w1_half]img1_path = os.path.join(target_path, f"{img_name[:-4]}_{i}.jpg")print("spilt img:", img1_path)cv2.imwrite(img1_path, img1)def split_img(img_file,wp = 2, hp = 2):img = cv2.imread(img_file)img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)img_h = img.shape[0] # 高度img_w = img.shape[1] # 宽度# h1_half = int(img_h / 2)# w1_half = int(img_w / 2)h1_half = img_h // hpw1_half = img_w // wpwh_sum = wp * hpprint(wh_sum)img_name = os.path.basename(img_file)for i in range(wh_sum):img1 = img[int(i / wp) * h1_half: h1_half * (int(i / wp) + 1), int(i % hp) * w1_half: (int(i % hp) + 1) * w1_half]img1_path = os.path.join(target_path, f"{img_name[:-wh_sum]}_{i}.jpg")print("spilt img:", img1_path)cv2.imwrite(img1_path, img1)if __name__ == '__main__':mkdir(target_path)for file in pathlib.Path(source_path).glob('**/*'):str = pathlib.Path(source_path)print(str)split_img(os.path.join(source_path, file),5,4)
4.3 训练模型
def train():"""训练模型:return:"""# model = YOLO("./ultralytics/cfg/models/v8/mtyolov8.yaml")model = YOLO("./yolov8n.pt")model.train(data="./ultralytics/cfg/datasets/coco8_HSG.yaml", epochs=400)result = model.val()
4.4 预测
import random
import time
from tqdm import tqdm
from ultralytics import YOLO
import torch
import cv2
import matplotlib.pyplot as plt
import os
import numpy as np
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'def load_train_model():"""加载训练好的模型:return:"""global modelmodel = YOLO("./best.pt")def load_image(image_path):"""加载图像:return: 图像"""print(image_path)if not os.path.exists(image_path):returnelse:return cv2.imread(image_path)def save_image(image):"""保存图像:return:"""# 创建目标文件夹中 images 和 labels 子文件夹os.makedirs(os.path.join("./", "predict_result"), exist_ok=True)# 移动图片到目标文件夹image_name = time.strftime('%Y-%m-%d-%H-%M-%S',time.localtime(time.time())) + ".jpg"image_path = "./predict_result/" + image_namecv2.imwrite(image_path, image)def get_images_path(image_path):"""获得文件加下所有图片的路径:return: 列表"""# 创建一个空列表来存储有效图片的路径images_path = []# 遍历 images 文件夹下的所有图片for image_name in os.listdir(image_path):# 获取图片的完整路径path = os.path.join(image_path, image_name)# 将图片路径添加到列表中images_path.append(path)return images_pathdef load_image_random(images_path):"""功能：在文件夹中随机加载一张图像:param image_path: 文件夹路径:return:"""pathname = random.choices(images_path)[0]image = load_image(pathname)return imagedef load_image_order(images_path):"""功能：在文件夹中按照先后顺序加载图片:param image_path: 文件夹路径:return:"""for pathname in images_path:image = load_image(pathname)yield imagedef train():"""训练模型:return:"""# model = YOLO("./ultralytics/cfg/models/v8/mtyolov8.yaml")model = YOLO("./yolov8n.pt")model.train(data="./ultralytics/cfg/datasets/coco8_HSG.yaml", epochs=400)result = model.val()def predict(pre_image):# 定义分割参数segment_width = 819segment_height = 750stride_w = 820 # w步长stride_h = 750 # h步长color_map = {'bright': (255, 0, 0),'acid': (0, 255, 0)# 其他类别的颜色可以根据需要添加}for y in range(0, pre_image.shape[0], stride_h):for x in range(0, pre_image.shape[1], stride_w):# 提取分割区域segment = pre_image[y:y + segment_height, x:x + segment_width]# 对分割区域进行目标检测results = model.predict(segment,conf=0.3)for result in results:boxes = result.boxesnames = result.namesnum = len(boxes.cls.cpu().numpy().astype(int))if num >= 1:for i in range(num):xyxy = boxes.xyxy.cpu().numpy().astype(int)[i]cls = boxes.cls.cpu().numpy().astype(int)[i]conf = boxes.conf.cpu().numpy()[i]color = color_map.get(names.get(cls), (0, 255, 0)) # 默认绿色# 将RGB格式的颜色转换为BGR格式color = (color[2], color[1], color[0])offset = 15x1 = xyxy[0] + x - offsety1 = xyxy[1] + y - offsetx2 = xyxy[2] + x + offsety2 = xyxy[3] + y + offsetcv2.rectangle(pre_image, (x1, y1), (x2, y2), color, 4)cv2.putText(pre_image, f"{names.get(cls)} {conf:.2f}", (x1, y1 - 15),cv2.FONT_HERSHEY_SIMPLEX, 1.5, color, 4)#cv2.imwrite("./result1.jpg",pre_image)#plt.title('Pre')#plt.imshow(cv2.cvtColor(pre_image, cv2.COLOR_BGR2RGB))#plt.show()return pre_image
5. 预测结果
- 测试集中明显缺陷已经可以检出，不明显的缺陷还需要继续优化。随着项目的推进，会有更多的缺陷图像补充进来，这样会极大提高模型的检测精度。
