Python Pytorch深度学习之数据加载和处理

#%%读取数据集
landmarks_frame=pd.read_csv('D:/Python/Pytorch/data/faces/face_landmarks.csv')
n=65
img_name=landmarks_frame.iloc[n,0]
landmarks=landmarks_frame.iloc[n,1:].values
landmarks=landmarks.astype('float').reshape(-1,2)
print('Image name :{}'.format(img_name))
print('Landmarks shape :{}'.format(landmarks.shape))
print('First 4 Landmarks:{}'.format(landmarks[:4]))

#%%编写显示人脸函数
def show_landmarks(image,landmarks):
    plt.imshow(image)
    plt.scatter(landmarks[:,0],landmarks[:,1],s=10,marker=".",c='r')
    plt.pause(0.001)
plt.figure()
show_landmarks(io.imread(os.path.join('D:/Python/Pytorch/data/faces/',img_name)),landmarks)
plt.show()

#%%数据集类——将数据集封装成一个类
class FaceLandmarksDataset(Dataset):
    def __init__(self,csv_file,root_dir,transform=None):
        # csv_file(string):待注释的csv文件的路径
        # root_dir(string):包含所有图像的目录
        # transform（callabele,optional）:一个样本上的可用的可选变换
        self.landmarks_frame=pd.read_csv(csv_file)
        self.root_dir=root_dir
        self.transform=transform
    def __len__(self):
        return len(self.landmarks_frame)
    def __getitem__(self, idx):
        img_name=os.path.join(self.root_dir,self.landmarks_frame.iloc[idx,0])
        image=io.imread(img_name)
        landmarks=self.landmarks_frame.iloc[idx,1:]
        landmarks=np.array([landmarks])
        landmarks=landmarks.astype('float').reshape(-1,2)
        sample={'image':image,'landmarks':landmarks}
        if self.transform:
            sample=self.transform(sample)
        return sample    

#%%数据可视化
# 将上面定义的类进行实例化并便利整个数据集
face_dataset=FaceLandmarksDataset(csv_file='D:/Python/Pytorch/data/faces/face_landmarks.csv', 
                                  root_dir='D:/Python/Pytorch/data/faces/')
fig=plt.figure()
for i in range(len(face_dataset)) :
    sample=face_dataset[i]
    print(i,sample['image'].shape,sample['landmarks'].shape)
    ax=plt.subplot(1,4,i+1)
    plt.tight_layout()
    ax.set_title('Sample #{}'.format(i))
    ax.axis('off')
    show_landmarks(**sample)
    if i==3:
        plt.show()
        break

# 将样本中的图像重新缩放到给定的大小
class Rescale(object):    
    def __init__(self,output_size):
        assert isinstance(output_size,(int,tuple))
        self.output_size=output_size
    #output_size 为int或tuple，如果是元组输出与output_size匹配，
    #如果是int，匹配较小的图像边缘到output_size保持纵横比相同
    def __call__(self,sample):
        image,landmarks=sample['image'],sample['landmarks']
        h,w=image.shape[:2]
        if isinstance(self.output_size, int):#输入参数是int
            if h>w:
                new_h,new_w=self.output_size*h/w,self.output_size
            else:
                new_h,new_w=self.output_size,self.output_size*w/h
        else:#输入参数是元组
            new_h,new_w=self.output_size
        new_h,new_w=int(new_h),int(new_w)
        img=transform.resize(image, (new_h,new_w))
        landmarks=landmarks*[new_w/w,new_h/h]
        return {'image':img,'landmarks':landmarks}

# 随机裁剪样本中的图像
class RandomCrop(object):
    def __init__(self,output_size):
        assert isinstance(output_size, (int,tuple))
        if isinstance(output_size, int):
            self.output_size=(output_size,output_size)
        else:
            assert len(output_size)==2
            self.output_size=output_size
    # 输入参数依旧表示想要裁剪后图像的尺寸，如果是元组其而包含两个元素直接复制长宽，如果是int，则裁剪为方形
    def __call__(self,sample):
        image,landmarks=sample['image'],sample['landmarks']
        h,w=image.shape[:2]
        new_h,new_w=self.output_size
        #确定图片裁剪位置
        top=np.random.randint(0,h-new_h)
        left=np.random.randint(0,w-new_w)
        image=image[top:top+new_h,left:left+new_w]
        landmarks=landmarks-[left,top]
        return {'image':image,'landmarks':landmarks}

#%%
# 将样本中的npdarray转换为Tensor
class ToTensor(object):
    def __call__(self,sample):
        image,landmarks=sample['image'],sample['landmarks']
        image=image.transpose((2,0,1))#交换颜色轴
        #numpy的图片是：Height*Width*Color
        #torch的图片是：Color*Height*Width
        return {'image':torch.from_numpy(image),
                'landmarks':torch.from_numpy(landmarks)}

#%%
scale=Rescale(256)
crop=RandomCrop(128)
composed=transforms.Compose([Rescale(256),RandomCrop(224)])    
# 在样本上应用上述变换
fig=plt.figure() 
sample=face_dataset[65]
for i,tsfrm in enumerate([scale,crop,composed]):
    transformed_sample=tsfrm(sample)
    ax=plt.subplot(1,3,i+1)
    plt.tight_layout()
    ax.set_title(type(tsfrm).__name__)
    show_landmarks(**transformed_sample)
plt.show()

transformed_dataset=FaceLandmarksDataset(csv_file='D:/Python/Pytorch/data/faces/face_landmarks.csv',
                                          root_dir='D:/Python/Pytorch/data/faces/',
                                          transform=transforms.Compose([
                                              Rescale(256),
                                              RandomCrop(224),
                                              ToTensor()
                                              ]))
for i in range(len(transformed_dataset)):
    sample=transformed_dataset[i]
    print(i,sample['image'].size(),sample['landmarks'].size())
    if i==3:
        break    

transformed_dataset=FaceLandmarksDataset(csv_file='D:/Python/Pytorch/data/faces/face_landmarks.csv',
                                          root_dir='D:/Python/Pytorch/data/faces/',
                                          transform=transforms.Compose([
                                              Rescale(256),
                                              RandomCrop(224),
                                              ToTensor()
                                              ]))
for i in range(len(transformed_dataset)):
    sample=transformed_dataset[i]
    print(i,sample['image'].size(),sample['landmarks'].size())
    if i==3:
        break