2020-12-21

文章：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 相关知识

MEGC2019 Github 代码