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

一、图片压接部位定位

，往往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/images

nc: 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 os

def 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, None

if __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 os

class YOLOApp:
    def __init__(self, root):
        self.root = root
        self.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_path
            self.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()

压接状态智能识别