耐张线夹压接图片智能识别

一、图片压接部位定位

，往往X射线照片是一个大图，进行图片压接部位定位目的是先找到需识别的部位，再进行识别时可排除其他图像部位的干扰，提高准确率。

1、图像准备

准备多个需进行压接状态智能识别的图片保存再source文件夹中

2、人工标注

使用labelImg 工具进行标注
图片文件夹设置为source
另存标注文件夹为annotations
保存格式为YOLO
快捷键：W为标注、D为下一张、A为上一张

3、训练

使用split_dataset.py 工具将图像与标注文件划分数据集

import os
import random
import shutil# 设置路径
source_dir = 'd:/Xradio/photo_enhance/source'
annotations_dir = 'd:/Xradio/photo_enhance/annotations'
output_dir = 'd:/Xradio/photo_enhance/dataset'# 创建输出目录
os.makedirs(os.path.join(output_dir, 'train/images'), exist_ok=True)
os.makedirs(os.path.join(output_dir, 'train/labels'), exist_ok=True)
os.makedirs(os.path.join(output_dir, 'val/images'), exist_ok=True)
os.makedirs(os.path.join(output_dir, 'val/labels'), exist_ok=True)
os.makedirs(os.path.join(output_dir, 'test/images'), exist_ok=True)
os.makedirs(os.path.join(output_dir, 'test/labels'), exist_ok=True)# 获取所有文件
image_files = [f for f in os.listdir(source_dir) if f.endswith(('.jpg', '.bmp'))]
random.shuffle(image_files)# 划分数据集
train_files = image_files[:int(len(image_files)*0.7)]
val_files = image_files[int(len(image_files)*0.7):int(len(image_files)*0.9)]
test_files = image_files[int(len(image_files)*0.9):]# 复制文件到相应目录
def copy_files(files, split):for f in files:# 复制图像shutil.copy(os.path.join(source_dir, f), os.path.join(output_dir, split, 'images', f))# 复制标注label_file = f.replace('.jpg', '.txt').replace('.bmp', '.txt')shutil.copy(os.path.join(annotations_dir, label_file),os.path.join(output_dir, split, 'labels', label_file))copy_files(train_files, 'train')
copy_files(val_files, 'val')
copy_files(test_files, 'test')print("数据集划分完成！")

创建labels.yaml 文件告诉训练系统各文件路径
其中train为训练数据、val为验证集、test为测试集

train: d:/Xradio/photo_enhance/dataset/train/images
val: d:/Xradio/photo_enhance/dataset/val/images
test: d:/Xradio/photo_enhance/dataset/test/imagesnc: 1
names: ['target']

使用以下指令进行训练：

yolo train data=labels.yaml model=yolov8n.pt epochs=50 imgsz=640 batch=8 amp=True device=0

训练结果保存在run/detect/train6/weights/best.pt

best.pt为训练的结果模型。

4、推理

编写推理程序，输入为模型与要检测的照片，输出为位置坐标

from ultralytics import YOLO
import cv2
import osdef main(image_path):# 加载YOLO模型model_path = 'runs/detect/train6/weights/best.pt'try:model = YOLO(model_path)except Exception as e:raise RuntimeError(f'无法加载模型: {str(e)}')# 验证图片路径if not os.path.exists(image_path):raise FileNotFoundError(f"文件不存在: {image_path}")# 读取图片original_image = cv2.imread(image_path)if original_image is None:raise ValueError("无法读取图片，请检查文件格式")# 进行推理results = model.predict(original_image)detections = results[0].boxes.xyxy.cpu().numpy()# 在图像上绘制检测框并返回第一个检测框的坐标annotated_image = original_image.copy()if len(detections) > 0:x1, y1, x2, y2 = map(int, detections[0])# 裁剪检测框内的图片部分cropped_image = original_image[y1:y2, x1:x2]return cropped_image, (x1, y1, x2, y2)else:return None, Noneif __name__ == '__main__':main()

5、UI界面

支持单张推理，也支持批量推理

import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from PIL import Image, ImageTk
from inference_app import main as run_inference
from ultralytics import YOLO
import cv2
import osclass YOLOApp:def __init__(self, root):self.root = rootself.root.title("YOLOv8 图像检测")self.root.geometry("800x600")# 创建界面元素self.create_widgets()def create_widgets(self):# 文件选择按钮self.select_button = tk.Button(self.root, text="选择图片", command=self.select_image)self.select_button.pack(pady=10)# 文件夹选择按钮self.select_folder_button = tk.Button(self.root, text="选择文件夹", command=self.select_folder)self.select_folder_button.pack(pady=10)# 图片显示区域self.image_label = tk.Label(self.root)self.image_label.pack(expand=True, fill=tk.BOTH)# 推理按钮self.infer_button = tk.Button(self.root, text="开始检测", command=self.run_detection, state=tk.DISABLED)self.infer_button.pack(pady=10)# 保存按钮self.save_button = tk.Button(self.root, text="保存结果", command=self.save_result, state=tk.DISABLED)self.save_button.pack(pady=10)def select_image(self):# 打开文件选择对话框file_path = filedialog.askopenfilename(filetypes=[("图片文件", "*.jpg *.jpeg *.png *.bmp")])if file_path:self.image_path = file_pathself.display_image(file_path)self.infer_button.config(state=tk.NORMAL)def display_image(self, file_path):# 显示原始图片image = Image.open(file_path)image.thumbnail((800, 600))self.photo = ImageTk.PhotoImage(image)self.image_label.config(image=self.photo)def run_detection(self):# 运行推理try:# 使用inference_app中的推理逻辑annotated_image, bbox = run_inference(self.image_path)# 显示带检测框的图片if annotated_image is not None:annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)annotated_image = Image.fromarray(annotated_image)annotated_image.thumbnail((800, 600))self.photo = ImageTk.PhotoImage(annotated_image)self.image_label.config(image=self.photo)self.save_button.config(state=tk.NORMAL)except Exception as e:messagebox.showerror("错误", f"推理失败: {str(e)}")def save_result(self):# 保存结果图片save_path = filedialog.asksaveasfilename(defaultextension=".jpg",filetypes=[("JPEG 文件", "*.jpg"), ("PNG 文件", "*.png")])if save_path:try:# 获取原始图片original_image = cv2.imread(self.image_path)# 运行推理获取检测框model_path = 'runs/detect/train6/weights/best.pt'model = YOLO(model_path)results = model.predict(original_image)detections = results[0].boxes.xyxy.cpu().numpy()if len(detections) > 0:x1, y1, x2, y2 = map(int, detections[0])# 裁剪图片cropped_image = original_image[y1:y2, x1:x2]# 保存裁剪后的图片cv2.imwrite(save_path, cropped_image)# 显示保存的图片self.display_image(save_path)messagebox.showinfo("成功", f"裁剪后的图片已保存到: {save_path}")else:messagebox.showwarning("警告", "未检测到目标，无法裁剪")except Exception as e:messagebox.showerror("错误", f"保存失败: {str(e)}")def select_folder(self):# 打开文件夹选择对话框folder_path = filedialog.askdirectory()if folder_path:# 创建输出目录output_dir = os.path.join(folder_path, "processed")os.makedirs(output_dir, exist_ok=True)# 获取所有图片文件image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))]# 创建进度条self.progress = tk.DoubleVar()self.progress_bar = ttk.Progressbar(self.root, variable=self.progress, maximum=len(image_files))self.progress_bar.pack(pady=10)# 批量处理for i, image_file in enumerate(image_files):try:# 更新进度self.progress.set(i + 1)self.root.update_idletasks()# 处理图片image_path = os.path.join(folder_path, image_file)annotated_image, _ = run_inference(image_path)# 保存结果output_path = os.path.join(output_dir, image_file)cv2.imwrite(output_path, annotated_image)except Exception as e:messagebox.showerror("错误", f"处理 {image_file} 失败: {str(e)}")continue# 处理完成messagebox.showinfo("完成", f"所有图片已处理完成，保存到: {output_dir}")self.progress_bar.pack_forget()if __name__ == "__main__":root = tk.Tk()app = YOLOApp(root)root.mainloop()

压接状态智能识别