文章:Revealing the Invisible with Model and Data Shrinking for Composite-database Micro-expression Recognition
一、代码
Datasets.py
import os
import torch
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
class MEGC2019(torch.utils.data.Dataset):
"""MEGC2019 dataset class with 3 categories"""
def __init__(self, imgList, transform=None):
self.imgPath = []
self.label = []
self.dbtype = []
with open(imgList,'r') as f:
for textline in f:
texts= textline.strip('\n').split(' ')
self.imgPath.append(texts[0])
self.label.append(int(texts[1]))
self.dbtype.append(int(texts[2]))
self.transform = transform
def __getitem__(self, idx):
img = Image.open("".join(self.imgPath[idx]),'r').convert('RGB')
# plt.imshow(img)
# plt.show()
if self.transform is not None:
img = self.transform(img)
return img, self.label[idx]
def __len__(self):
return len(self.imgPath)
class MEGC2019_SI(torch.utils.data.Dataset):
"""MEGC2019_SI dataset class with 3 categories and other side information"""
def __init__(self, imgList, transform=None):
self.imgPath = []
self.label = []
self.dbtype = []
with open(imgList,'r') as f:
for textline in f:
texts= textline.strip('\n').split(' ')
self.imgPath.append(texts[0])
self.label.append(int(texts[1]))
self.dbtype.append(int(texts[2]))
self.transform = transform
def __getitem__(self, idx):
img = Image.open("".join(self.imgPath[idx]),'r').convert('RGB')
# plt.imshow(img)
# plt.show()
if self.transform is not None:
img = self.transform(img)
return {"data":img, "class_label":self.label[idx], 'db_label':self.dbtype[idx]}
def __len__(self):
return len(self.imgPath)
class MEGC2019_FOLDER(torch.utils.data.Dataset):
"""MEGC2019 dataset class with 3 categories, organized in folders"""
def __init__(self, rootDir, transform=None):
labels = os.listdir(rootDir)
labels.sort()
self.fileList = []
self.label = []
self.imgPath = []
for subfolder in labels:
label = []
imgPath = []
files = os.listdir(os.path.join(rootDir, subfolder))
files.sort()
self.fileList.extend(files)
label = [int(subfolder) for file in files]
imgPath = [os.path.join(rootDir, subfolder,file) for file in files]
self.label.extend(label)
self.imgPath.extend(imgPath)
self.transform = transform
def __getitem__(self, idx):
img = Image.open(self.imgPath[idx],'r').convert('RGB')
# plt.imshow(img)
# plt.show()
if self.transform is not None:
img = self.transform(img)
return {"data":img, "class_label":self.label[idx]}
def __len__(self):
return len(self.fileList)
二、代码理解
继承 torch.utils.data.Dataset
:
这是自定义数据集方法的抽象类;
继承这个类可以定义自己的数据集类;
继承这个类只需要覆写__len__和__getitem__这两个方法;
len(self) 是实例使用len()方法时调用;
getitem(self)是实例读取指定索引的元素时调用;
torch.utils.data.Dataset 相关知识