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 os
def 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:0
if 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 pathlib
source_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 True
else:
print(path + ' 目录已存在')
return False
def 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 // 2
w1_half = img_w // 2
img_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 // hp
w1_half = img_w // wp
wh_sum = wp * hp
print(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 model
model = YOLO("./best.pt")
def load_image(image_path):
"""
加载图像
:return: 图像
"""
print(image_path)
if not os.path.exists(image_path):
return
else:
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_name
cv2.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_path
def load_image_random(images_path):
"""
功能：在文件夹中随机加载一张图像
:param image_path: 文件夹路径
:return:
"""
pathname = random.choices(images_path)[0]
image = load_image(pathname)
return image
def load_image_order(images_path):
"""
功能：在文件夹中按照先后顺序加载图片
:param image_path: 文件夹路径
:return:
"""
for pathname in images_path:
image = load_image(pathname)
yield image
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()
def predict(pre_image):
# 定义分割参数
segment_width = 819
segment_height = 750
stride_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.boxes
names = result.names
num = 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 = 15
x1 = xyxy[0] + x - offset
y1 = xyxy[1] + y - offset
x2 = xyxy[2] + x + offset
y2 = xyxy[3] + y + offset
cv2.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. 预测结果
- 测试集中明显缺陷已经可以检出，不明显的缺陷还需要继续优化。随着项目的推进，会有更多的缺陷图像补充进来，这样会极大提高模型的检测精度。
