Python 随机划分训练和验证集保存成pkl

最新推荐文章于 2024-07-14 10:29:57 发布

Along1617188

最新推荐文章于 2024-07-14 10:29:57 发布

阅读量641

点赞数

分类专栏： Python Numpy 医学图像处理文章标签： python pkl

本文链接：https://blog.youkuaiyun.com/Along1617188/article/details/124339086

版权

Python 同时被 3 个专栏收录

77 篇文章

订阅专栏

医学图像处理

15 篇文章

订阅专栏

Numpy

8 篇文章

订阅专栏

此篇博客介绍了如何使用Python获取目录下的nii文件，并将图像和标签列表按8:2的比例随机划分成训练集和验证集，然后将结果分别保存为pkl文件。通过get_train_val_list函数读取文件，write_pkl函数实现数据划分和存储。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import os
import pickle
import numpy as np

def get_train_val_list(work_path):
    '''
    todo:   get train val list 
    '''
    image_list = []
    mask_list = []
    case_list = os.listdir(work_path)   
    case_list.sort()
    for case in case_list:
        case_path = os.path.join(work_path, case)
        origin_case_path = os.path.join(case_path, 'origin.nii.gz') 
        mask_case_path = os.path.join(case_path, 'label.nii.gz')  
        image_list.append(origin_case_path)
        mask_list.append(mask_case_path)
    return image_list, mask_list

def write_pkl(save_path, image_list, mask_list):
    '''
    todo:   write train_list and val_list as 8:2 to pkl
    '''
    with open(os.path.join(save_path, 'train.pkl'), "wb") as train_pkl, open(os.path.join(save_path, 'val.pkl'), "wb") as val_pkl:
        data = dict()
        for i in range(len(image_list)):
            data['image_path'] = image_list[i]
            data['mask_path'] = mask_list[i]

            # split train and val data
            if np.random.random()<0.8:
                pickle.dump(data, train_pkl)
            else:
                pickle.dump(data, val_pkl)

def main():
    root_path = '/data/xx/nii'
    save_path = '/home/yy'
    image_list, mask_list = get_train_val_list(root_path)
    write_pkl(save_path,image_list, mask_list)
if __name__=="__main__":
    main()