三维重构终版
This commit is contained in:
		
							
								
								
									
										23
									
								
								yolov8/predict.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										23
									
								
								yolov8/predict.py
									
									
									
									
									
										Normal file
									
								
							@@ -0,0 +1,23 @@
 | 
			
		||||
from ultralytics import YOLO
 | 
			
		||||
import os
 | 
			
		||||
 | 
			
		||||
# 确保 if __name__ == '__main__': 结构,这是一个好习惯
 | 
			
		||||
if __name__ == '__main__':
 | 
			
		||||
    # 1. 加载你训练好的最佳模型
 | 
			
		||||
    # !! 修改为你自己的 best.pt 路径 !!
 | 
			
		||||
    model_path = r'runs_up/detect/train/weights/best.pt'
 | 
			
		||||
    model = YOLO(model_path)
 | 
			
		||||
 | 
			
		||||
    # 2. 指定你要预测的图片或文件夹
 | 
			
		||||
    # 可以是单张图片路径,也可以是整个文件夹的路径
 | 
			
		||||
    # 强烈建议使用验证集里的图片,或者一些全新的测试图片
 | 
			
		||||
    source_path = r'train_data_up/images/val' # 预测整个验证集文件夹
 | 
			
		||||
 | 
			
		||||
    # 3. 执行预测
 | 
			
		||||
    # save=True:  会将画好框的图片保存下来
 | 
			
		||||
    # conf=0.5:   只显示置信度大于 0.5 的预测结果,可以调整这个值
 | 
			
		||||
    results = model.predict(source=source_path, save=True, conf=0.5)
 | 
			
		||||
 | 
			
		||||
    # 预测结果会默认保存在 runs_up/detect/predictX 文件夹下
 | 
			
		||||
    print("\n预测完成!")
 | 
			
		||||
    # 你可以从 results 对象中获取详细信息,但对于可视化验证,直接去看保存的图片更方便。
 | 
			
		||||
							
								
								
									
										116
									
								
								yolov8/split_dataset.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										116
									
								
								yolov8/split_dataset.py
									
									
									
									
									
										Normal file
									
								
							@@ -0,0 +1,116 @@
 | 
			
		||||
import os
 | 
			
		||||
import random
 | 
			
		||||
import shutil
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
def split_dataset(image_dir, label_dir, output_dir, split_ratio=0.8):
 | 
			
		||||
    """
 | 
			
		||||
    自动将图片和标签划分为训练集和验证集。
 | 
			
		||||
 | 
			
		||||
    Args:
 | 
			
		||||
        image_dir (str): 原始图片文件夹路径。
 | 
			
		||||
        label_dir (str): 转换后的 YOLO 标签 (.txt) 文件夹路径。
 | 
			
		||||
        output_dir (str): 整理好的数据集输出根目录 (例如 'weld_dataset')。
 | 
			
		||||
        split_ratio (float): 训练集所占的比例,例如 0.8 代表 80% 训练, 20% 验证。
 | 
			
		||||
    """
 | 
			
		||||
    print("开始划分数据集...")
 | 
			
		||||
 | 
			
		||||
    # --- 1. 路径设置和文件夹创建 ---
 | 
			
		||||
    train_img_path = os.path.join(output_dir, 'images', 'train')
 | 
			
		||||
    val_img_path = os.path.join(output_dir, 'images', 'val')
 | 
			
		||||
    train_label_path = os.path.join(output_dir, 'labels', 'train')
 | 
			
		||||
    val_label_path = os.path.join(output_dir, 'labels', 'val')
 | 
			
		||||
 | 
			
		||||
    # 创建所有必要的文件夹
 | 
			
		||||
    os.makedirs(train_img_path, exist_ok=True)
 | 
			
		||||
    os.makedirs(val_img_path, exist_ok=True)
 | 
			
		||||
    os.makedirs(train_label_path, exist_ok=True)
 | 
			
		||||
    os.makedirs(val_label_path, exist_ok=True)
 | 
			
		||||
 | 
			
		||||
    # --- 2. 文件匹配 ---
 | 
			
		||||
    # 获取所有标签文件的基础名(不含扩展名)
 | 
			
		||||
    label_files = [os.path.splitext(f)[0] for f in os.listdir(label_dir) if f.endswith('.txt')]
 | 
			
		||||
 | 
			
		||||
    # 查找对应的图片文件(支持多种格式)
 | 
			
		||||
    image_files_map = {}
 | 
			
		||||
    supported_formats = ['.jpg', '.jpeg', '.png', '.bmp']
 | 
			
		||||
    for f in os.listdir(image_dir):
 | 
			
		||||
        base_name, ext = os.path.splitext(f)
 | 
			
		||||
        if ext.lower() in supported_formats:
 | 
			
		||||
            image_files_map[base_name] = f
 | 
			
		||||
 | 
			
		||||
    # 找出图片和标签都存在的文件对
 | 
			
		||||
    valid_files = [base_name for base_name in label_files if base_name in image_files_map]
 | 
			
		||||
 | 
			
		||||
    if not valid_files:
 | 
			
		||||
        print(f"错误:在图片目录 '{image_dir}' 和标签目录 '{label_dir}' 之间未找到任何匹配的文件对。")
 | 
			
		||||
        print("请确保图片和标签的文件名(除扩展名外)完全一致。")
 | 
			
		||||
        return
 | 
			
		||||
 | 
			
		||||
    print(f"共找到 {len(valid_files)} 个有效的图片-标签对。")
 | 
			
		||||
 | 
			
		||||
    # --- 3. 随机划分 ---
 | 
			
		||||
    random.shuffle(valid_files)
 | 
			
		||||
    split_index = int(len(valid_files) * split_ratio)
 | 
			
		||||
    train_files = valid_files[:split_index]
 | 
			
		||||
    val_files = valid_files[split_index:]
 | 
			
		||||
 | 
			
		||||
    # --- 4. 复制文件到目标位置 ---
 | 
			
		||||
    def copy_files(file_list, dest_img_path, dest_label_path):
 | 
			
		||||
        for base_name in file_list:
 | 
			
		||||
            # 复制图片
 | 
			
		||||
            img_name = image_files_map[base_name]
 | 
			
		||||
            shutil.copy(os.path.join(image_dir, img_name), dest_img_path)
 | 
			
		||||
            # 复制标签
 | 
			
		||||
            label_name = base_name + '.txt'
 | 
			
		||||
            shutil.copy(os.path.join(label_dir, label_name), dest_label_path)
 | 
			
		||||
 | 
			
		||||
    print(f"正在复制 {len(train_files)} 个文件到训练集...")
 | 
			
		||||
    copy_files(train_files, train_img_path, train_label_path)
 | 
			
		||||
 | 
			
		||||
    print(f"正在复制 {len(val_files)} 个文件到验证集...")
 | 
			
		||||
    copy_files(val_files, val_img_path, val_label_path)
 | 
			
		||||
 | 
			
		||||
    print("\n数据集划分完成!")
 | 
			
		||||
    print(f"训练集: {len(train_files)} 张图片 | 验证集: {len(val_files)} 张图片")
 | 
			
		||||
    print(f"数据已整理至 '{output_dir}' 文件夹。")
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
if __name__ == '__main__':
 | 
			
		||||
    # --- 用户需要配置的参数 ---
 | 
			
		||||
 | 
			
		||||
    # # 1. 原始图片文件夹路径
 | 
			
		||||
    # # !! 重要 !!: 请将这里的路径修改为您实际存放图片的文件夹
 | 
			
		||||
    # # 可能是 'faster-rcnn/JPEGImages' 或其他名称
 | 
			
		||||
    # ORIGINAL_IMAGE_DIR = '../OpenCV/data_bottom/test3'
 | 
			
		||||
    #
 | 
			
		||||
    # # 2. 转换后的 YOLO 标签文件夹路径
 | 
			
		||||
    # YOLO_LABEL_DIR = 'data_bottom'
 | 
			
		||||
    #
 | 
			
		||||
    # # 3. 最终输出的数据集文件夹
 | 
			
		||||
    # OUTPUT_DATASET_DIR = 'train_data'
 | 
			
		||||
    #
 | 
			
		||||
    # # 4. 训练集比例 (0.8 表示 80% 训练, 20% 验证)
 | 
			
		||||
    # SPLIT_RATIO = 0.9
 | 
			
		||||
    #
 | 
			
		||||
    # # --- 运行主函数 ---
 | 
			
		||||
    # split_dataset(ORIGINAL_IMAGE_DIR, YOLO_LABEL_DIR, OUTPUT_DATASET_DIR, SPLIT_RATIO)
 | 
			
		||||
    #
 | 
			
		||||
    # --- 用户需要配置的参数 ---
 | 
			
		||||
 | 
			
		||||
    # 1. 原始图片文件夹路径
 | 
			
		||||
    # !! 重要 !!: 请将这里的路径修改为您实际存放图片的文件夹
 | 
			
		||||
    # 可能是 'faster-rcnn/JPEGImages' 或其他名称
 | 
			
		||||
    ORIGINAL_IMAGE_DIR = '../label/up'
 | 
			
		||||
 | 
			
		||||
    # 2. 转换后的 YOLO 标签文件夹路径
 | 
			
		||||
    YOLO_LABEL_DIR = 'data_up'
 | 
			
		||||
 | 
			
		||||
    # 3. 最终输出的数据集文件夹
 | 
			
		||||
    OUTPUT_DATASET_DIR = 'train_data_up'
 | 
			
		||||
 | 
			
		||||
    # 4. 训练集比例 (0.8 表示 80% 训练, 20% 验证)
 | 
			
		||||
    SPLIT_RATIO = 0.9
 | 
			
		||||
 | 
			
		||||
    # --- 运行主函数 ---
 | 
			
		||||
    split_dataset(ORIGINAL_IMAGE_DIR, YOLO_LABEL_DIR, OUTPUT_DATASET_DIR, SPLIT_RATIO)
 | 
			
		||||
							
								
								
									
										9
									
								
								yolov8/train_data/weld.yaml
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										9
									
								
								yolov8/train_data/weld.yaml
									
									
									
									
									
										Normal file
									
								
							@@ -0,0 +1,9 @@
 | 
			
		||||
# 数据集路径设置
 | 
			
		||||
# '.' 表示相对于yolo命令运行的目录,这里假设你会从 innovate_project 根目录运行
 | 
			
		||||
path: ../train_data  # 数据集根目录
 | 
			
		||||
train: images/train  # 训练图片文件夹 (相对于 path)
 | 
			
		||||
val: images/val      # 验证图片文件夹 (相对于 path)
 | 
			
		||||
 | 
			
		||||
# 类别信息
 | 
			
		||||
nc: 1  # number of classes: 类别数量
 | 
			
		||||
names: ['Space weld workpiece']  # class names: 类别名称列表
 | 
			
		||||
							
								
								
									
										20
									
								
								yolov8/train_yolo.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										20
									
								
								yolov8/train_yolo.py
									
									
									
									
									
										Normal file
									
								
							@@ -0,0 +1,20 @@
 | 
			
		||||
from ultralytics import YOLO
 | 
			
		||||
 | 
			
		||||
# 将所有执行代码都放入这个 if 语句块中
 | 
			
		||||
if __name__ == '__main__':
 | 
			
		||||
    # 1. 加载一个预训练模型
 | 
			
		||||
    # 这行代码可以在 if 语句内外,但为了代码清晰和规范,建议放入
 | 
			
		||||
    model = YOLO('yolov8n.pt')
 | 
			
		||||
 | 
			
		||||
    # 2. 训练模型
 | 
			
		||||
    # 核心的训练启动代码必须在 if 语句内
 | 
			
		||||
    results = model.train(data='train_data_up/weld.yaml',
 | 
			
		||||
                          epochs=50,
 | 
			
		||||
                          imgsz=640,
 | 
			
		||||
                          device=0,
 | 
			
		||||
                          workers=8)  # 可以显式指定 workers 数量
 | 
			
		||||
 | 
			
		||||
    # 3. (可选) 打印训练结果保存的路径
 | 
			
		||||
    print("训练完成!")
 | 
			
		||||
    # 注意: 在新版ultralytics中,results可能不直接包含save_dir,但训练日志会打印出来
 | 
			
		||||
    # 训练结果通常保存在 runs_bottom/detect/trainX 目录下
 | 
			
		||||
							
								
								
									
										119
									
								
								yolov8/voc_to_yolo.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										119
									
								
								yolov8/voc_to_yolo.py
									
									
									
									
									
										Normal file
									
								
							@@ -0,0 +1,119 @@
 | 
			
		||||
import xml.etree.ElementTree as ET
 | 
			
		||||
import os
 | 
			
		||||
import glob
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
def voc_to_yolo(xml_file_path, output_dir, class_mapping):
 | 
			
		||||
    """
 | 
			
		||||
    将单个 PASCAL VOC anntation (.xml) 文件转换为 YOLO (.txt) 格式。
 | 
			
		||||
 | 
			
		||||
    Args:
 | 
			
		||||
        xml_file_path (str): 输入的 .xml 文件路径。
 | 
			
		||||
        output_dir (str): 输出 .txt 文件的目标文件夹。
 | 
			
		||||
        class_mapping (dict): 类别名称到类别ID的映射字典。
 | 
			
		||||
    """
 | 
			
		||||
    try:
 | 
			
		||||
        # 解析 XML 文件
 | 
			
		||||
        tree = ET.parse(xml_file_path)
 | 
			
		||||
        root = tree.getroot()
 | 
			
		||||
 | 
			
		||||
        # 获取图像尺寸
 | 
			
		||||
        size = root.find('size')
 | 
			
		||||
        if size is None:
 | 
			
		||||
            print(f"警告: 在 {xml_file_path} 中未找到 <size> 标签,跳过此文件。")
 | 
			
		||||
            return
 | 
			
		||||
 | 
			
		||||
        img_width = int(size.find('width').text)
 | 
			
		||||
        img_height = int(size.find('height').text)
 | 
			
		||||
 | 
			
		||||
        # 准备用于写入的YOLO标注列表
 | 
			
		||||
        yolo_annotations = []
 | 
			
		||||
 | 
			
		||||
        # 遍历所有 object
 | 
			
		||||
        for obj in root.findall('object'):
 | 
			
		||||
            # 获取类别名称
 | 
			
		||||
            class_name = obj.find('name').text
 | 
			
		||||
            if class_name not in class_mapping:
 | 
			
		||||
                print(f"警告: 类别 '{class_name}' 不在预定义的 class_mapping 中,跳过此物体。")
 | 
			
		||||
                continue
 | 
			
		||||
 | 
			
		||||
            class_id = class_mapping[class_name]
 | 
			
		||||
 | 
			
		||||
            # 获取边界框坐标
 | 
			
		||||
            bndbox = obj.find('bndbox')
 | 
			
		||||
            xmin = float(bndbox.find('xmin').text)
 | 
			
		||||
            ymin = float(bndbox.find('ymin').text)
 | 
			
		||||
            xmax = float(bndbox.find('xmax').text)
 | 
			
		||||
            ymax = float(bndbox.find('ymax').text)
 | 
			
		||||
 | 
			
		||||
            # --- 核心转换公式 ---
 | 
			
		||||
            x_center = (xmin + xmax) / 2.0 / img_width
 | 
			
		||||
            y_center = (ymin + ymax) / 2.0 / img_height
 | 
			
		||||
            width = (xmax - xmin) / img_width
 | 
			
		||||
            height = (ymax - ymin) / img_height
 | 
			
		||||
 | 
			
		||||
            # 将结果添加到列表
 | 
			
		||||
            yolo_annotations.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
 | 
			
		||||
 | 
			
		||||
        # 如果文件中有有效的物体,则写入 .txt 文件
 | 
			
		||||
        if yolo_annotations:
 | 
			
		||||
            # 构建输出文件名
 | 
			
		||||
            base_filename = os.path.basename(xml_file_path)
 | 
			
		||||
            txt_filename = os.path.splitext(base_filename)[0] + '.txt'
 | 
			
		||||
            output_path = os.path.join(output_dir, txt_filename)
 | 
			
		||||
 | 
			
		||||
            # 写入文件
 | 
			
		||||
            with open(output_path, 'w') as f:
 | 
			
		||||
                f.write('\n'.join(yolo_annotations))
 | 
			
		||||
 | 
			
		||||
            # print(f"成功转换: {xml_file_path} -> {output_path}")
 | 
			
		||||
 | 
			
		||||
    except Exception as e:
 | 
			
		||||
        print(f"处理文件 {xml_file_path} 时发生错误: {e}")
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
def main():
 | 
			
		||||
    # --- 用户需要配置的参数 ---
 | 
			
		||||
 | 
			
		||||
    # 1. 定义你的类别和对应的ID (从0开始)
 | 
			
		||||
    # 根据你的截图,你只有一个类别 "Space weld workpiece"
 | 
			
		||||
    # 请确保这里的名称与你XML文件中的<name>标签完全一致!
 | 
			
		||||
    CLASS_MAPPING = {
 | 
			
		||||
        'Space weld workpiece': 0,
 | 
			
		||||
        # 如果有其他类别,在这里添加,例如:
 | 
			
		||||
        # 'weld_seam': 1,
 | 
			
		||||
    }
 | 
			
		||||
 | 
			
		||||
    # 2. 定义输入和输出文件夹
 | 
			
		||||
    # 输入文件夹: 存放 .xml 文件的目录
 | 
			
		||||
    input_xml_dir = '../label/up_xml'
 | 
			
		||||
 | 
			
		||||
    # 输出文件夹: 存放转换后的 .txt 文件的目录
 | 
			
		||||
    output_txt_dir = 'data_up'
 | 
			
		||||
 | 
			
		||||
    # --- 脚本执行部分 ---
 | 
			
		||||
 | 
			
		||||
    # 自动创建输出文件夹(如果不存在)
 | 
			
		||||
    if not os.path.exists(output_txt_dir):
 | 
			
		||||
        os.makedirs(output_txt_dir)
 | 
			
		||||
        print(f"已创建输出文件夹: {output_txt_dir}")
 | 
			
		||||
 | 
			
		||||
    # 查找所有 .xml 文件
 | 
			
		||||
    xml_files = glob.glob(os.path.join(input_xml_dir, '*.xml'))
 | 
			
		||||
 | 
			
		||||
    if not xml_files:
 | 
			
		||||
        print(f"错误: 在目录 '{input_xml_dir}' 中没有找到任何 .xml 文件。请检查路径。")
 | 
			
		||||
        return
 | 
			
		||||
 | 
			
		||||
    print(f"找到 {len(xml_files)} 个 .xml 文件。开始转换...")
 | 
			
		||||
 | 
			
		||||
    # 遍历并转换每个文件
 | 
			
		||||
    for xml_file in xml_files:
 | 
			
		||||
        voc_to_yolo(xml_file, output_txt_dir, CLASS_MAPPING)
 | 
			
		||||
 | 
			
		||||
    print("\n转换完成!")
 | 
			
		||||
    print(f"所有 YOLO 格式的标签文件已保存在: {output_txt_dir}")
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
if __name__ == '__main__':
 | 
			
		||||
    main()
 | 
			
		||||
		Reference in New Issue
	
	Block a user