TransUnet官方代码训练自己数据集（彩色RGB3通道图像的分割） |
您所在的位置：网站首页 › 图像分割csdn › TransUnet官方代码训练自己数据集（彩色RGB3通道图像的分割） |
TransUnet官方代码训练自己数据集（彩色RGB3通道图像的分割）
|
***************************************************
码字不易，收藏之余，别忘了给我点个赞吧！
***************************************************
---------Start
官方代码：https://github.com/Beckschen/TransUNet
目的：训练5个类别的汽车部件分割任务（测试在另一篇博客中）
CSDN数据集免费下载
实现效果：
转化数据集的代码如下，会将images中的图像和labels中的标签生成一个.npz文件。 def npz(): #图像路径 path = r'G:\dataset\car-segmentation\train\images\*.png' #项目中存放训练所用的npz文件路径 path2 = r'G:\dataset\Unet\TransUnet-ori\data\Synapse\train_npz\\' for i,img_path in enumerate(glob.glob(path)): #读入图像 image = cv2.imread(img_path) image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB) #读入标签 label_path = img_path.replace('images','labels') label = cv2.imread(label_path,flags=0) #保存npz np.savez(path2+str(i),image=image,label=label) print('------------',i) # 加载npz文件 # data = np.load(r'G:\dataset\Unet\Swin-Unet-ori\data\Synapse\train_npz\0.npz', allow_pickle=True) # image, label = data['image'], data['label'] print('ok')生成的文件在 data\Synapse\train_npz文件夹中，如下图，也可以自己定义生成的路径，然后把文件复制到data\Synapse\train_npz文件中。
官方下载地址 CSDN下载地址[推荐] 进入网站后，点击imagenet21k文件夹。 找到datasets/dataset_synapse.py文件中的Synapse_dataset类，修改__getitem__函数。 def __getitem__(self, idx): if self.split == "train": slice_name = self.sample_list[idx].strip('\n') data_path = self.data_dir+"/"+slice_name+'.npz' data = np.load(data_path) image, label = data['image'], data['label'] else: slice_name = self.sample_list[idx].strip('\n') data_path = self.data_dir+"/"+slice_name+'.npz' data = np.load(data_path) image, label = data['image'], data['label'] image = torch.from_numpy(image.astype(np.float32)) image = image.permute(2,0,1) label = torch.from_numpy(label.astype(np.float32)) sample = {'image': image, 'label': label} if self.transform: sample = self.transform(sample) sample['case_name'] = self.sample_list[idx].strip('\n') return sample找到datasets/dataset_synapse.py文件中的RandomGenerator类，修改__call__函数。 def __call__(self, sample): image, label = sample['image'], sample['label'] if random.random() > 0.5: image, label = random_rot_flip(image, label) elif random.random() > 0.5: image, label = random_rotate(image, label) x, y,_ = image.shape if x != self.output_size[0] or y != self.output_size[1]: image = zoom(image, (self.output_size[0] / x, self.output_size[1] / y,1), order=3) # why not 3? label = zoom(label, (self.output_size[0] / x, self.output_size[1] / y), order=0) image = torch.from_numpy(image.astype(np.float32)) image = image.permute(2,0,1) label = torch.from_numpy(label.astype(np.float32)) sample = {'image': image, 'label': label.long()} return sample至此，数据读取的部分已经修改完毕！ 5. 配置训练参数认真检查各个参数是否正确，这里的路径都是 ‘./’(当前目录下)，不是"…/"，训练时，batch_size通常大于1，我这里设置有误！类别数可根据你的任务定！
设置trainer.py文件中的DataLoader函数中的num_workers=0 至此，所有代码修改完毕！ 总结：以上修改内容针对彩色图像的分割任务， 由于仅文字表述某些操作存在局限性，故只能简略应答，有任何问题可下方留言评论。 |
官方数据集格式，data文件夹，Synapse文件夹，test_vol_h5文件夹，train_npz文件夹手动创建！ 
data\Synapse\train_npz文件夹中存放的是训练集样本，按照同样的方式生成测试集样本，存放在data\Synapse\test_vol_h5文件夹中。 
我的训练集203个样本，测试集3个样本。npz文件生成完成之后，找到train.txt和test_vol.txt，手动将文件里面的内容清空，split_data.py这个文件直接无视。自己写一个函数读取train_npz中所有的文件名称，然后将文件名称写入train.txt文件，一个名称一行，如下图所示。同理可完成test_vol.txt文件制作。 
至此，数据集制作完毕！！！代码会先去train.txt文件中读取训练样本的名称，然后根据名称再去train_npz文件夹下读取npz文件。所以每一步都很重要，必须正确！
下载这个权重文件即可。 
手动创建如下多个文件夹，存放刚刚下载完毕的权重，注意名称跟我的保持一致！ 
至此，预训练权重已下载完毕。
图片大小设置，越大越耗显存。 
【本文地址】
今日新闻 |
推荐新闻 |