Celyn_blog

步骤一：构建YOLO数据集

由于比赛原始数据集较大，我们采样部分数据构建训练集和验证集：

if os.path.exists('yolo_seg_dataset'):
    shutil.rmtree('yolo_seg_dataset')

os.makedirs('yolo_seg_dataset/train')
os.makedirs('yolo_seg_dataset/valid')

def normalize_polygon(polygon, img_width, img_height):
    return [(x / img_width, y / img_height) for x, y in polygon]

# 采样训练集
for row in training_anno.iloc[:10000].iterrows():
    shutil.copy(row[1].Path, 'yolo_seg_dataset/train')

    img = cv2.imread(row[1].Path)
    img_height, img_width = img.shape[:2]
    txt_filename = os.path.join('yolo_seg_dataset/train/' + row[1].Path.split('/')[-1][:-4] + '.txt')
    with open(txt_filename, 'w') as up:
        for polygon in row[1].Polygons:
            normalized_polygon = normalize_polygon(polygon, img_width, img_height)
            normalized_coords = ' '.join([f'{coord[0]:.3f} {coord[1]:.3f}' for coord in normalized_polygon])
            up.write(f'0 {normalized_coords}\n')
 
 # 采用验证集     
 for row in training_anno.iloc[10000:10150].iterrows():
    shutil.copy(row[1].Path, 'yolo_seg_dataset/valid')

    img = cv2.imread(row[1].Path)
    img_height, img_width = img.shape[:2]
    txt_filename = os.path.join('yolo_seg_dataset/valid/' + row[1].Path.split('/')[-1][:-4] + '.txt')
    with open(txt_filename, 'w') as up:
        for polygon in row[1].Polygons:
            normalized_polygon = normalize_polygon(polygon, img_width, img_height)
            normalized_coords = ' '.join([f'{coord[0]:.3f} {coord[1]:.3f}' for coord in normalized_polygon])
            up.write(f'0 {normalized_coords}\n')

逐行代码分析

好的，我们逐行分析这段代码，它的主要功能是将训练集和验证集的数据从原始数据集中复制到新的目录中，并将多边形的坐标标准化为相对坐标。

代码逐行分析

if os.path.exists('yolo_seg_dataset'):
    shutil.rmtree('yolo_seg_dataset')

• 功能: 检查yolo_seg_dataset目录是否存在。如果存在，则使用shutil.rmtree删除该目录及其内容。shutil常用于文件与目录的处理。

• 目的: 确保每次运行代码时，数据集是干净的，没有旧的数据。

os.makedirs('yolo_seg_dataset/train')
os.makedirs('yolo_seg_dataset/valid')

• 功能: 创建两个新的子目录train和valid，用于存放训练集和验证集数据。

• 目的: 为之后的文件复制和标签文件创建所需的目录结构。