點擊上方藍字關注我們

微信公眾號：OpenCV學堂

關注獲取更多計算機視覺與深度學習知識

大家好，這是輕松學Pytorch系列的第六篇分享，本篇你將學會如何從頭開始制作自己的數據集，并通過DataLoader實現加載。本文以人臉Landmard五點的數據集標定與之制作為例來說明pytorch中如何實現自定義數據集讀取與加載。

數據來源

首先要實現人臉landmark五點的數據標定，就得找到人臉數據，我使用的人臉數據是celebA數據集，大概有20W張多點，我從中選擇了1000張，然后通過OpenCV寫了個程序對人臉進行了簡單的裁剪。然后還選擇了一個音樂MV(上次就被人打call的宇少)通過opencv實現采集了一些人臉數據，這個數據的好處是有不同的光照，各種角度，豐富了數據的多樣性。這些數據加起來1500張左右。圖示如下：

Landmark標定

我這里選擇的對得到1500張圖像做數據標注，剛開始的選擇標定工具都讓我頭疼，這個是我第一次標定一系列的點，經過一番嘗試之后，終于發現一個很好用的工具，同時支持人臉檢測與五點標定。貼上地址：

https://github.com/Mukosame/Face-Annotation-Tool

廢話也不多說了，只說一句話，簡單靠譜，然后我就對這個1500張圖像進行五點標定，本來我可以不這么干的，我可以用其它的模型來直接找這些圖像的landmark五點然后生成文件即可，但是我還是決定手動標注一番。結果讓我眼睛疼了兩天之后，終于給標注好拉，發誓以后再也不干這種活了，我太難了。截圖如下：

現在自定義數據已經準備完畢，下面就應該是pytorch登場了。

自定義數據集實現

基于Pytorch中的torch.utils.data.Dataset類實現自定義的FaceLandmarksDataset類，主要是重寫了getitem這個方法。完整的代碼實現如下：

1class?FaceLandmarksDataset(Dataset):
2????def?__init__(self,?txt_file):
3????????self.transform?=?transforms.Compose([transforms.ToTensor()])
4????????lines?=?[]
5????????with?open(txt_file)?as?read_file:
6????????????for?line?in?read_file:
7????????????????line?=?line.replace('\n',?'')
8????????????????lines.append(line)
9????????self.landmarks_frame?=?lines
10
11????def?__len__(self):
12????????return?len(self.landmarks_frame)
13
14????def?num_of_samples(self):
15????????return?len(self.landmarks_frame)
16
17????def?__getitem__(self,?idx):
18????????if?torch.is_tensor(idx):
19????????????idx?=?idx.tolist()
20????????contents?=?self.landmarks_frame[idx].split('\t')
21????????image_path?=?contents[0]
22????????img?=?cv.imread(image_path)??#?BGR?order
23????????h,?w,?c?=?img.shape
24????????#?rescale
25????????img?=?cv.resize(img,?(64,?64))
26????????img?=?(np.float32(img)?/255.0?-?0.5)?/?0.5
27????????landmarks?=?np.zeros(10,?dtype=np.float32)
28????????for?i?in?range(1,?len(contents),?2):
29????????????landmarks[i?-?1]?=?np.float32(contents[i])?/?w
30????????????landmarks[i]?=?np.float32(contents[i?+?1])?/?h
31????????landmarks?=?landmarks.astype('float32').reshape(-1,?2)
32????????#?H,?W?C?to?C,?H,?W
33????????img?=?img.transpose((2,?0,?1))
34????????sample?=?{'image':?torch.from_numpy(img),?'landmarks':?torch.from_numpy(landmarks)}
35????????return?sample

加載與顯示

實現了自定義的Dataset類之后，就可以通過自定義的Dataset來構建一個DataLoader對象實現數據的加載跟批次處理，對自定義的dataset完成測試。代碼如下：

1ds?=?FaceLandmarksDataset("D:/facedb/Face-Annotation-Tool/landmark_output.txt")
2for?i?in?range(len(ds)):
3????sample?=?ds[i]
4????print(i,?sample['image'].size(),?sample['landmarks'].size())
5????if?i?==?3:
6????????break
7
8dataloader?=?DataLoader(ds,?batch_size=4,?shuffle=True,?num_workers=4)
9#?data?loader
10for?i_batch,?sample_batched?in?enumerate(dataloader):
11????print(i_batch,?sample_batched['image'].size(),?sample_batched['landmarks'].size())

運行顯示如下：

關注我們，后臺輸入關鍵字 landmark 獲取本人辛苦標注的數據集。

記得點贊支持，這是本人繼續寫下去的動力！

?推薦閱讀?

輕松學Pytorch–環境搭建與基本語法

Pytorch輕松學-構建淺層神經網絡

輕松學pytorch-構建卷積神經網絡

輕松學Pytorch –構建循環神經網絡

輕松學Pytorch-使用卷積神經網絡實現圖像分類

不積硅步，無以至千里

總結

以上是生活随笔為你收集整理的数据集制作_轻松学Pytorch自定义数据集制作与使用的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。