以讀取VOC2012語義分割數(shù)據(jù)集為例，具體見代碼注釋：
VocDataset.py
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
|
from PIL import Imageimport torchimport torch.utils.data as dataimport numpy as npimport osimport torchvisionimport torchvision.transforms as transformsimport time#VOC數(shù)據(jù)集分類對(duì)應(yīng)顏色標(biāo)簽VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0], [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128], [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0], [0, 64, 128]]#顏色標(biāo)簽空間轉(zhuǎn)到序號(hào)標(biāo)簽空間，就他媽這里浪費(fèi)巨量的時(shí)間,這里還他媽的有問題def voc_label_indices(colormap, colormap2label): """Assign label indices for Pascal VOC2012 Dataset.""" idx = ((colormap[:, :, 2] * 256 + colormap[ :, :,1]) * 256+ colormap[:, :,0]) #out = np.empty(idx.shape, dtype = np.int64) out = colormap2label[idx] out=out.astype(np.int64)#數(shù)據(jù)類型轉(zhuǎn)換 end = time.time() return outclass MyDataset(data.Dataset):#創(chuàng)建自定義的數(shù)據(jù)讀取類 def __init__(self, root, is_train, crop_size=(320,480)): self.rgb_mean =(0.485, 0.456, 0.406) self.rgb_std = (0.229, 0.224, 0.225) self.root=root self.crop_size=crop_size images = []#創(chuàng)建空列表存文件名稱 txt_fname = '%s/ImageSets/Segmentation/%s' % (root, 'train.txt' if is_train else 'val.txt') with open(txt_fname, 'r') as f: self.images = f.read().split() #數(shù)據(jù)名稱整理 self.files = [] for name in self.images: img_file = os.path.join(self.root, "JPEGImages/%s.jpg" % name) label_file = os.path.join(self.root, "SegmentationClass/%s.png" % name) self.files.append({ "img": img_file, "label": label_file, "name": name }) self.colormap2label = np.zeros(256**3) #整個(gè)循環(huán)的意思就是將顏色標(biāo)簽映射為單通道的數(shù)組索引 for i, cm in enumerate(VOC_COLORMAP): self.colormap2label[(cm[2] * 256 + cm[1]) * 256 + cm[0]] = i #按照索引讀取每個(gè)元素的具體內(nèi)容 def __getitem__(self, index): datafiles = self.files[index] name = datafiles["name"] image = Image.open(datafiles["img"]) label = Image.open(datafiles["label"]).convert('RGB')#打開的是PNG格式的圖片要轉(zhuǎn)到rgb的格式下，不然結(jié)果會(huì)比較要命 #以圖像中心為中心截取固定大小圖像，小于固定大小的圖像則自動(dòng)填0 imgCenterCrop = transforms.Compose([ transforms.CenterCrop(self.crop_size), transforms.ToTensor(), transforms.Normalize(self.rgb_mean, self.rgb_std),#圖像數(shù)據(jù)正則化 ]) labelCenterCrop = transforms.CenterCrop(self.crop_size) cropImage=imgCenterCrop(image) croplabel=labelCenterCrop(label) croplabel=torch.from_numpy(np.array(croplabel)).long()#把標(biāo)簽數(shù)據(jù)類型轉(zhuǎn)為torch #將顏色標(biāo)簽圖轉(zhuǎn)為序號(hào)標(biāo)簽圖 mylabel=voc_label_indices(croplabel, self.colormap2label) return cropImage,mylabel #返回圖像數(shù)據(jù)長度 def __len__(self): return len(self.files) |
Train.py
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
|
import matplotlib.pyplot as pltimport torch.utils.data as dataimport torchvision.transforms as transformsimport numpy as npfrom PIL import Imagefrom VocDataset import MyDataset#VOC數(shù)據(jù)集分類對(duì)應(yīng)顏色標(biāo)簽VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0], [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128], [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0], [0, 64, 128]]root='../data/VOCdevkit/VOC2012'train_data=MyDataset(root,True)trainloader = data.DataLoader(train_data, 4)#從數(shù)據(jù)集中拿出一個(gè)批次的數(shù)據(jù)for i, data in enumerate(trainloader): getimgs, labels= data img = transforms.ToPILImage()(getimgs[0]) labels = labels.numpy()#tensor轉(zhuǎn)numpy labels=labels[0]#獲得批次標(biāo)簽集中的一張標(biāo)簽圖像 labels = labels.transpose((1,0))#數(shù)組維度切換，將第1維換到第0維，第0維換到第1維 ##將單通道索引標(biāo)簽圖片映射回顏色標(biāo)簽圖片 newIm= Image.new('RGB', (480, 320))#創(chuàng)建一張與標(biāo)簽大小相同的圖片，用以顯示標(biāo)簽所對(duì)應(yīng)的顏色 for i in range(0, 480): for j in range(0, 320): sele=labels[i][j]#取得坐標(biāo)點(diǎn)對(duì)應(yīng)像素的值 newIm.putpixel((i, j), (int(VOC_COLORMAP[sele][0]), int(VOC_COLORMAP[sele][1]), int(VOC_COLORMAP[sele][2]))) #顯示圖像和標(biāo)簽 plt.figure("image") ax1 = plt.subplot(1,2,1) ax2 = plt.subplot(1,2,2) plt.sca(ax1) plt.imshow(img) plt.sca(ax2) plt.imshow(newIm) plt.show() |
以上這篇Pytorch 實(shí)現(xiàn)數(shù)據(jù)集自定義讀取就是小編分享給大家的全部內(nèi)容了，希望能給大家一個(gè)參考，也希望大家多多支持服務(wù)器之家。
原文鏈接：https://blog.csdn.net/weixinhum/article/details/86232734
