目标检测——GDXray数据集转为YOLO格式

news2025/1/23 15:03:40

关于该数据集的介绍可以看我写的另一篇博客：链接

论文题目：《GDXray: The Database of X-ray Images for Nondestructive Testing》
论文链接：https://link.springer.com/article/10.1007/s10921-015-0315-7

Github链接： https://github.com/computervision-xray-testing/GDXray/blob/main/README.md

以GDXray的Baggages为例，其原数据集下载下来如下所示：

其中每个文件夹里都有ground_truth.txt，即坐标框相关的信息，ground_truth.txt内容如下：

知道了每个数字对应的含义，接下来的任务就简单了，因此只需要针对相同的ID创建同样的TXT的Label，并且对应图片的ID名即可，参考代码如下：
注：该代码仅是针对一组图片的，所有的图片都需转换的话再套一个循环即可！

import os


def map_number_to_filename(dir_name, number):
    """
    将数字映射为特定的文件名格式。

    参数:
    - number: 一个浮点数，表示映射的键。

    返回:
    - str: 按照给定格式映射后的文件名。
    """


    # 将浮点数转换为字符串，并去掉不必要的小数点和之后的所有0
    number_str = format(number, '.0f')
    
    # 将字符串格式化为至少4位的数字，不足部分前面补0
    number_str_padded = number_str.zfill(4)
    
    # 构造文件名
    img_name = f"{dir_name}_{number_str_padded}.png"
    label_name = f"{dir_name}_{number_str_padded}.txt"

    
    return label_name, img_name


import cv2

def get_image_dimensions(image_path, pos_list):
    """
    使用OpenCV获取指定图片文件的宽度和高度。
    
    参数:
    - image_path: 图片文件的路径。
    
    返回:
    - (width, height): 图片的宽度和高度的元组。
    """
    templist = []
    # 使用cv2.imread读取图片，0表示以原始格式读取
    image = cv2.imread(image_path, 0)
    
    # 检查图片是否正确读取
    if image is None:
        raise FileNotFoundError(f"The image at path {image_path} could not be found or is not a valid image file.")
    
    # 图片的尺寸可以通过image.shape属性获取，它返回一个三元组(height, width, channels)
    height, width = image.shape[:2]

    b_w = float(pos_list[1]) - float(pos_list[0])
    b_h = float(pos_list[3]) - float(pos_list[2])
    b_x = (float(pos_list[0]) + float(pos_list[1]))/2.0
    b_y = (float(pos_list[2]) + float(pos_list[3]))/2.0

    x = float(b_x)/width
    y = float(b_y)/height
    w = float(b_w)/width
    h = float(b_h)/height
    
    templist.append(x)
    templist.append(y)
    templist.append(w)
    templist.append(h)
    return templist


if __name__ == "__main__":
    dir_path = r"F:\BaiduNetdiskDownload\DataSet\GDXray\Baggages\Baggages\B0001"
    dir_basename = os.path.basename(dir_path)
    ground_truth_txt = os.path.join(dir_path,'ground_truth.txt')
    save_txt_dir = os.path.join(os.path.dirname(dir_path), f"YOLO_{dir_basename}")
    if not os.path.exists(save_txt_dir):
        os.makedirs(save_txt_dir)

    f = open(ground_truth_txt,'r',encoding='utf-8')
    contents = f.readlines()

    for content in contents:
        content = content.strip().split("   ")
        filename, img_name = map_number_to_filename(dir_name=dir_basename, number=float(content[0]))
        img_path = os.path.join(dir_path, img_name)
        yolo_pos_list = get_image_dimensions(img_path, pos_list=content[1:])

        save_filename = os.path.join(save_txt_dir,filename)
        f_s = open(save_filename,'a',encoding='utf-8')
        f_s.write("0 ")
        f_s.write(str(yolo_pos_list[0]))
        f_s.write(" ")
        f_s.write(str(yolo_pos_list[1]))
        f_s.write(" ")
        f_s.write(str(yolo_pos_list[2]))
        f_s.write(" ")
        f_s.write(str(yolo_pos_list[3]))

拿一组数据集来进行实验，B0001文件夹