輕松學(xué)Pytorch-自定義數(shù)據(jù)集制作與使用

大家好，這是輕松學(xué)Pytorch系列的第六篇分享，本篇你將學(xué)會(huì)如何從頭開始制作自己的數(shù)據(jù)集，并通過(guò)DataLoader實(shí)現(xiàn)加載。本文以人臉Landmard五點(diǎn)的數(shù)據(jù)集標(biāo)定與之制作為例來(lái)說(shuō)明pytorch中如何實(shí)現(xiàn)自定義數(shù)據(jù)集讀取與加載。

首先要實(shí)現(xiàn)人臉landmark五點(diǎn)的數(shù)據(jù)標(biāo)定，就得找到人臉數(shù)據(jù)，我使用的人臉數(shù)據(jù)是celebA數(shù)據(jù)集，大概有20W張多點(diǎn)，我從中選擇了1000張，然后通過(guò)OpenCV寫了個(gè)程序?qū)θ四樳M(jìn)行了簡(jiǎn)單的裁剪。然后還選擇了一個(gè)音樂(lè)MV(上次就被人打call的宇少)通過(guò)opencv實(shí)現(xiàn)采集了一些人臉數(shù)據(jù)，這個(gè)數(shù)據(jù)的好處是有不同的光照，各種角度，豐富了數(shù)據(jù)的多樣性。這些數(shù)據(jù)加起來(lái)1500張左右。圖示如下：

我這里選擇的對(duì)得到1500張圖像做數(shù)據(jù)標(biāo)注，剛開始的選擇標(biāo)定工具都讓我頭疼，這個(gè)是我第一次標(biāo)定一系列的點(diǎn)，經(jīng)過(guò)一番嘗試之后，終于發(fā)現(xiàn)一個(gè)很好用的工具，同時(shí)支持人臉檢測(cè)與五點(diǎn)標(biāo)定。貼上地址：

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

廢話也不多說(shuō)了，只說(shuō)一句話，簡(jiǎn)單靠譜，然后我就對(duì)這個(gè)1500張圖像進(jìn)行五點(diǎn)標(biāo)定，本來(lái)我可以不這么干的，我可以用其它的模型來(lái)直接找這些圖像的landmark五點(diǎn)然后生成文件即可，但是我還是決定手動(dòng)標(biāo)注一番。結(jié)果讓我眼睛疼了兩天之后，終于給標(biāo)注好拉，發(fā)誓以后再也不干這種活了，我太難了。截圖如下：

現(xiàn)在自定義數(shù)據(jù)已經(jīng)準(zhǔn)備完畢，下面就應(yīng)該是pytorch登場(chǎng)了。

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

 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

實(shí)現(xiàn)了自定義的Dataset類之后，就可以通過(guò)自定義的Dataset來(lái)構(gòu)建一個(gè)DataLoader對(duì)象實(shí)現(xiàn)數(shù)據(jù)的加載跟批次處理，對(duì)自定義的dataset完成測(cè)試。代碼如下：

 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())

運(yùn)行顯示如下：

關(guān)注我們，后臺(tái)輸入關(guān)鍵字 landmark 獲取本人辛苦標(biāo)注的數(shù)據(jù)集。

記得點(diǎn)贊支持，這是本人繼續(xù)寫下去的動(dòng)力！

下載1：OpenCV-Contrib擴(kuò)展模塊中文版教程
在「小白學(xué)視覺」公眾號(hào)后臺(tái)回復(fù)：擴(kuò)展模塊中文教程，即可下載全網(wǎng)第一份OpenCV擴(kuò)展模塊教程中文版，涵蓋擴(kuò)展模塊安裝、SFM算法、立體視覺、目標(biāo)跟蹤、生物視覺、超分辨率處理等二十多章內(nèi)容。
下載2：Python視覺實(shí)戰(zhàn)項(xiàng)目52講
在「小白學(xué)視覺」公眾號(hào)后臺(tái)回復(fù)：Python視覺實(shí)戰(zhàn)項(xiàng)目，即可下載包括圖像分割、口罩檢測(cè)、車道線檢測(cè)、車輛計(jì)數(shù)、添加眼線、車牌識(shí)別、字符識(shí)別、情緒檢測(cè)、文本內(nèi)容提取、面部識(shí)別等31個(gè)視覺實(shí)戰(zhàn)項(xiàng)目，助力快速學(xué)校計(jì)算機(jī)視覺。
下載3：OpenCV實(shí)戰(zhàn)項(xiàng)目20講
在「小白學(xué)視覺」公眾號(hào)后臺(tái)回復(fù)：OpenCV實(shí)戰(zhàn)項(xiàng)目20講，即可下載含有20個(gè)基于OpenCV實(shí)現(xiàn)20個(gè)實(shí)戰(zhàn)項(xiàng)目，實(shí)現(xiàn)OpenCV學(xué)習(xí)進(jìn)階。

交流群

歡迎加入公眾號(hào)讀者群一起和同行交流，目前有SLAM、三維視覺、傳感器、自動(dòng)駕駛、計(jì)算攝影、檢測(cè)、分割、識(shí)別、醫(yī)學(xué)影像、GAN、算法競(jìng)賽等微信群（以后會(huì)逐漸細(xì)分），請(qǐng)掃描下面微信號(hào)加群，備注：”昵稱+學(xué)校/公司+研究方向“，例如：”張三 + 上海交大 + 視覺SLAM“。請(qǐng)按照格式備注，否則不予通過(guò)。添加成功后會(huì)根據(jù)研究方向邀請(qǐng)進(jìn)入相關(guān)微信群。請(qǐng)勿在群內(nèi)發(fā)送廣告，否則會(huì)請(qǐng)出群，謝謝理解~