dataset.py
目录:
- 前言
- 观察数据
- 书写代码
- 函数解释
前言
在步骤中需要写自己的dataset类,并将label和image一一对应后返回。
观察数据
在书写dataset前最重要的就是要观察数据集,对数据集进行分析,比如了解图片大小,通道数目,他的ndarray的dtype类型等等。甚至可以自己书写一个脚本,对数据本身进行分析。
如下以果蝇电镜图为例,我们观察数据集,知道Images中存放30张原图,Labels中存放30张已经分割好的图片,每张图片以.png
方式进行存储;Images和Labels中的图片通过文件名进行一一对应,如Labels中0.png
对应Images中0.png
,进行共有30组对应数据;根据model.py
文件了解需要传入的Tensor类型,思考如何将现有数据转为需要的Tensor返回。
知道以上信息后,我们书写dataset.py将image和label一一对应返回。
书写代码
在本步骤中,我们需要告诉程序如何读入你的数据,并且做一些预处理。我们的DIYDataset类继承自Dataset类,并重写__init__
,__len__
和__getitem__
三个魔法方法。网络训练权重为float32,所以传入数据也一般要为float32。
__init__
方法向外索取3个输入数据:
- 读取数据路径
- 是训练集还是验证集(因为你训练集和验证集往往是在两个不同的文件夹)
- 你使用的预处理方法(以transform为主,transform也可以根据验证集还是测试集调用不同的trans)
__getitem__
需要返回image和对应label的Tensor__len__
用来返回集合中图片个数
以S-BIAD25为例,代码如下:
'''
该类用于返回dataset
input和label都是[512, 512]的图片,需要将其转换为[3, 512, 512]才能transforms
'''
# --- add path and DIY package
import sys, os
root_path = os.path.dirname(os.path.dirname(__file__))
project_path = os.path.dirname(__file__)
sys.path.append(root_path)
sys.path.append(project_path)
from _utils import tensor_info # 耦合了
# ---
from torch.utils.data.dataset import Dataset
from PIL import Image
from utils.utils import cvtColor,resize_image
from torchvision import transforms
class CellDataset(Dataset):
"""返回细胞分隔的dataset"""
def __init__(self, path:str, transforms:object=None) -> None:
super().__init__()
self.path = path
self.labels = os.listdir(os.path.join(path,'Labels'))
self.transforms = transforms
def __len__(self):
return len(self.labels)
def __getitem__(self, index):
yield_name = self.labels[index]
# 返回label图片
label_path = os.path.join(self.path, 'Labels', yield_name)
label_image = Image.open(label_path)
label_image = cvtColor(label_image) # 将单通道灰度图片, 四通道png转换为三通道RGB图片。因为此处torchvision.transforms只接受三通道图片,pytorch中模型一般只能训练dtype=float32的Tensor。
label_image = resize_image(label_image,(512, 512))[0] # 裁剪图片大小为512 * 512
# 返回训练图片
image_path = os.path.join(self.path, 'Images', yield_name)
image = Image.open(image_path)
image = cvtColor(image)
image = resize_image(image,(512, 512))[0]
# 返回单张input, label
return self.transforms(image), self.transforms(label_image)
class MouseDataset(Dataset):
"""返回细胞分隔的dataset"""
def __init__(self, path) -> None:
super().__init__()
self.path = path
self.f_actin = os.path.join(path, "F-actin")
self.labels_factin = os.listdir(self.f_actin)
# self.labels = os.listdir(os.path.join(path,'Labels'))
self.transforms = transforms.Compose([transforms.ToTensor()])
def __len__(self):
return len(self.labels_factin)
def __getitem__(self, index):
yield_name = self.labels_factin[index][-19:]
# 返回label图片
label_path = os.path.join(self.path, 'F-actin', "img_568"+yield_name)
label_image = Image.open(label_path)
label_image = cvtColor(label_image) # 将单通道灰度图片, 四通道png转换为三通道RGB图片。因为此处torchvision.transforms只接受三通道图片,pytorch中模型一般只能训练dtype=float32的Tensor。
label_image = resize_image(label_image,(512, 512))[0] # 裁剪图片大小为512 * 512
# 返回训练图片
image_path = os.path.join(self.path, 'retardance', "img_Retardance"+yield_name)
image = Image.open(image_path)
image = cvtColor(image)
image = resize_image(image,(512, 512))[0]
# 返回单张input, label
return self.transforms(image), self.transforms(label_image)
if __name__ == "__main__":
"""test"""
data_path = "/home/yingmuzhi/_data/S-BIAD25"
cell_dataset = MouseDataset(data_path)
input, label = cell_dataset[0] # 同 input, label = cell_dataset.__getitem__(0)
tensor_info.TensorInfo(input).show_info()
tensor_info.TensorInfo(label).show_info()
测试结果
我们在if __name__ == "__main__":
中进行测试,结果如下:
函数解释
对一些常用函数进行解释,可以当做字典查看
torchvision.transforms.ToTensor
只接受PIL.Image类型的对象或者numpy.ndarray类型的对象,将上面两个对象转为torch.Tensor对象。函数定义如下:
参考链接https://pytorch.org/vision/stable/generated/torchvision.transforms.ToTensor.html?highlight=totensor#torchvision.transforms.ToTensor
__init__
在该方法中需要根据传入的path
参数和train
参数找到你的测试集或者训练集的物理地址,并将集合中的images和labels的物理地址存储在list中,以供后面方法使用。案例如下:
def __init__(self, root: str, train: bool, transforms: object=None):
super(DriveDataset, self).__init__()
self.flag = "training" if train else "test" # 由 train: bool 的布尔值来判断是取train还是test
data_root = os.path.join(root, "DRIVE", self.flag)
assert os.path.exists(data_root), f"path '{data_root}' does not exists."
self.transforms = transforms
img_names = [i for i in os.listdir(os.path.join(data_root, "images")) if i.endswith(".tif")]
self.img_list = [os.path.join(data_root, "images", i) for i in img_names] # 返回所有images地址的list
self.manual = [os.path.join(data_root, "1st_manual", i.split("_")[0] + "_manual1.gif") # 返回所有manuals地址的list
for i in img_names]
# check manual files
for i in self.manual:
if os.path.exists(i) is False:
raise FileNotFoundError(f"file {i} does not exists.")
self.roi_mask = [os.path.join(data_root, "mask", i.split("_")[0] + f"_{self.flag}_mask.gif") # 返回所有mask地址的list
for i in img_names]
# check mask files
for i in self.roi_mask:
if os.path.exists(i) is False:
raise FileNotFoundError(f"file {i} does not exists.")
__len__
返回测试集或者验证集中Image数量,因为Image和Label往往是一一对应的,所以返回哪个其实都一样。三维数据可能存在多对一情况。
案例如下:
def __len__(self):
return len(self.img_list)
__getitem__
该方法根据image和label的物理地址,用PIL
打开图片,再用transforms
处理Image返回Tensor,最后返回处理过的Tensor类型元组(image, label)。
在该方法中,你可以使用PIL
处理图像(mode),也可以将PIL
转为numpy
使用numpy处理图片(元素类型dtype),也可以使用Transforms
处理图片(Normalization)等。
__getitem__
案例见下:
def __getitem__(self, idx):
"""将Image转为RGB, 将label转为L"""
img = Image.open(self.img_list[idx]).convert('RGB')
manual = Image.open(self.manual[idx])
manual = manual.convert('L')
manual = np.array(manual) / 255
roi_mask = Image.open(self.roi_mask[idx]).convert('L')
roi_mask = 255 - np.array(roi_mask)
# 将manual图片和Imae进行处理
mask = np.clip(manual + roi_mask, a_min=0, a_max=255)
# 这里转回PIL的原因是,transforms中是对PIL数据进行处理
mask = Image.fromarray(mask)
if self.transforms is not None:
img, mask = self.transforms(img, mask)
return img, mask
使用PIL处理
参看链接https://blog.csdn.net/qq_43369406/article/details/127781871
使用Numpy处理
参看链接https://blog.csdn.net/qq_43369406/article/details/127781871
使用transforms处理
参看链接[coming soon]
这段代码我们常写在train.py
中。在进行transforms累加时候,我们常将所需要的transforms
全部添加至一个list中,再将这个list给transforms.Compose
掉,注意一定要添加transforms.ToTensor
方法。transform的更多内容可以参考笔者的transforms博客。
在调用transforms时完整逻辑如下:
# 获取dataset
train_dataset = DriveDataset(args.data_path,
train=True,
transforms=get_transform(train=True, mean=mean, std=std))
# 获取tranforms
def get_transform(train, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
"""对要获取的transforms进行判断,看是测试集的dataset还是验证集"""
base_size = 565
crop_size = 480
if train:
return SegmentationPresetTrain(base_size, crop_size, mean=mean, std=std)
else:
return SegmentationPresetEval(mean=mean, std=std)
# 测试集transforms
class SegmentationPresetTrain:
def __init__(self, base_size, crop_size, hflip_prob=0.5, vflip_prob=0.5,
mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
min_size = int(0.5 * base_size)
max_size = int(1.2 * base_size)
trans = [T.RandomResize(min_size, max_size)]
if hflip_prob > 0:
trans.append(T.RandomHorizontalFlip(hflip_prob))
if vflip_prob > 0:
trans.append(T.RandomVerticalFlip(vflip_prob))
trans.extend([
T.RandomCrop(crop_size),
T.ToTensor(),
T.Normalize(mean=mean, std=std),
])
self.transforms = T.Compose(trans)
def __call__(self, img, target):
return self.transforms(img, target)
增加了魔法方法__call__
只是为了对image/input,label/target一块进行transforms。
torchvision.transforms.Compose()和torchvision.transforms.[functions]
Compose()类可以将多个transforms对象合在一起给数据进行预处理,常以多个transforms对象的list形式传入Compose()中,如transforms.Compose([])
。transforms.[functions]则可以对多个输入数据进行变换,Compose()函数原型如下:
# 原型
torchvision.transforms.Compose(transforms)
# example
data_transform = {
"train": transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]),
"val": transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
}
torchvision.datasets.ImageFolder()类
ImageFolder()用于加载数据集,它说到底还是继承自torch.utils.data.Dataset
,后续可以作为Dataset对象直接传入torch.utils.data.DataLoader
中。对于ImageFolder(),root是要加载数据集的路径,transforms是对数据进行预处理的方式,函数原型和example如下:
# 原型
torchvision.datasets.ImageFolder(root: str, transform: Optional[Callable] = None,
target_transform: Optional[Callable] = None, loader: Callable[[str], Any] = <function default_loader>, is_valid_file: Optional[Callable[[str], bool]] = None)
# example
train_dataset = datasets.ImageFolder(root=project_path + "/flower_data/train", transform=data_transform["train"])
python 列表生成式和生成器
python 常见列表生成式 和 生成器,其中列表生成式以[]圈起,生成器以()圈起,常见列表生成器如下:
而生成器则是把中括号改成小括号。
# 列表生成式
generate_list = [i for i in range(10) if i < 5]
# 生成器
cla_dict = dict((val, key) for key, val in flower_list.items()) # 解包取出键值对,生成dict字典,赋值给cla_dict