【Pytorch】Face Landmarks

# http://pytorch123.com/ThirdSection/LearningPyTorch/
import os
import torch
import pandas as pd
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader,Dataset
from torchvision import transforms, utils

# 忽略警告
import warnings
warnings.filterwarnings("ignore")

# Part I
landmarks_frame = pd.read_csv("/home/DataSet/faces/face_landmarks.csv")
#print(landmarks_frame) # [69 rows x 137 columns] the 1 col is name, 2-137 are features

n = 65
img_name = landmarks_frame.iloc[n,0] # get the 65 number image's name
print("***",img_name)

landmarks = landmarks_frame.iloc[n,1:].as_matrix() # get the 65 number image's 136 features

landmarks = landmarks.astype('float').reshape(-1,2) # (136,) to (68, 2)  convenient to draw in axis

# draw a data
def show_landmarks(image,landmarks):
    plt.imshow(image)
    plt.scatter(landmarks[:,0],landmarks[:,1],s=10,marker='.',c='red')

show_landmarks(io.imread(os.path.join('/home/DataSet/faces/',img_name)),landmarks)

# Part II
# build a dataset class,{'image',image,'landmarks':landmarks}
class FaceLandmarkDataset(Dataset):
    def __init__(self,csv_file,root_dir,transform=None):
        """
        :param csv_file: string
        :param root_dir: string
        :param tranform: data augmentation
        """
        self.landmarks_frame = pd.read_csv(csv_file) # origin annotation
        self.root_dir = root_dir # the content of images
        self.transform = transform # data augmentation

    def __len__(self):
        return len(self.landmarks_frame)

    def __getitem__(self, idx): # according idx to load image
        img_name =  os.path.join(self.root_dir,
                                 self.landmarks_frame.iloc[idx,0]) # get picture's name
        image = io.imread(img_name) # read a image
        landmarks = self.landmarks_frame.iloc[idx,1:] # get features
        landmarks = np.array([landmarks])
        landmarks = landmarks.astype('float').reshape(-1,2) # change to 2 dim for localization
        sample = {'image':image,'landmarks':landmarks} # save as a dictionary

        if self.transform: # data augmentation
            sample = self.transform(sample)

        return sample

# show some examples by using FaceLandmarkDataset class
face_dataset = FaceLandmarkDataset(csv_file="/home/DataSet/faces/face_landmarks.csv",
                                   root_dir="/home/DataSet/faces/")

#print(face_dataset[0])

fig = plt.figure() # another figure window

print("Num, Image_Shape, Landmarks_Shape")
for i in range(len(face_dataset)):
    sample = face_dataset[i] # here __getitem__ is activel

    print(i,sample['image'].shape,sample['landmarks'].shape)

    ax = plt.subplot(1,4,i+1)
    ax.axis('off')
    show_landmarks(sample['image'],sample['landmarks'])
    #show_landmarks(**sample) # amazing

    if i == 3: # only show four images
        break
#plt.show()

# Part III
# transform rescale randomcrop and

class Rescale(object):
    """
    将样本中的图像重新缩放到给定大小。.
    Args:
        output_size（tuple或int）：所需的输出大小。 如果是元组，则输出为
        与output_size匹配。 如果是int，则匹配较小的图像边缘到output_size保持纵横比相同。
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))  # occurs when the condition is false
        self.output_size = output_size

    def __call__(self, sample):
        image,landmarks = sample['image'],sample['landmarks']

        h,w = image.shape[:2]
        if isinstance(self.output_size,int):
            if h > w:
                new_h,new_w = self.output_size * h/w, self.output_size # equal the short edge and keep the ratio
            else:
                new_h,new_w = self.output_size, self.output_size * w/h # equal the short edge and keep the ratio
        else:
            new_h, new_w = self.output_size # output_size is tuple

        new_h, new_w = int(new_h),int(new_w) # chang to int

        img = transform.resize(image,(new_h,new_w))

        # h and w are swapped for landmarks because for images,
        # x and y axes are axis 1 and 0 respectively
        landmarks = landmarks * [new_w / w, new_h /h] # origin is [w,h], now is [new_w,new_h]

        return {'image':img,'landmarks':landmarks}

class RandomCrop(object):
    """随机裁剪样本中的图像.

    Args:
       output_size（tuple或int）：所需的输出大小。 如果是int，方形裁剪是。
    """
    def __init__(self,output_size):
        assert isinstance(output_size, (int, tuple))  # occurs when the condition is false
        if isinstance(output_size,int):
            self.output_size = (output_size,output_size)
        else:
            assert len(output_size) == 2
            self.output_size = output_size

    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) # 返回随机整数，范围区间为[low,high)
        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}

class ToTensor(object):
    """将样本中的ndarrays转换为Tensors."""
    def __call__(self,sample):
        image, landmarks = sample['image'], sample['landmarks']
        # numpy is H*W*C
        # tensor in Torch is C*H*W
        image = image.transpose((2,0,1))
        return {'image': torch.from_numpy(image), 'landmarks': torch.from_numpy(landmarks)}

scale = Rescale(256)  # short edge to 256
crop = RandomCrop(128) # crop 224 from 256
composed = transforms.Compose([Rescale(256),RandomCrop(224)]) # scale and crop together

fig = plt.figure()

sample = face_dataset[65] # (160,160,3)
# print(sample["image"].shape)

for i,trans in enumerate([scale,crop,composed]):
    transformed_sample = trans(sample)

    ax = plt.subplot(1,3,i+1)
    plt.tight_layout()
    ax.set_title(type(trans).__name__)
    show_landmarks(**transformed_sample)

# part IV
transformed_dataset = FaceLandmarkDataset(csv_file="/home/DataSet/faces/face_landmarks.csv",
                                          root_dir="/home/DataSet/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
# 批量处理数据 * 打乱数据 * 使用多线程 multi processing worker 并行加载数据, instead of 'for' loop to load data

dataloader = DataLoader(transformed_dataset,
                        batch_size=4,
                        shuffle=True,
                        num_workers=4)

# 辅助功能：显示批次

def show_landmarks_batch(sample_batch):
    """Show image with landmarks for a batch of samples."""
    image_batch,landmarks_batch = sample_batch['image'],sample_batch['landmarks']
    batch_size = len(image_batch)
    im_size = image_batch.size(2)
    grid_border_size = 2
    grid = utils.make_grid(image_batch) # make_grid的作用是将若干幅图像拼成一幅图像
    plt.imshow(grid.numpy().transpose((1,2,0)))

    for i in range(batch_size):
        plt.scatter(landmarks_batch[i,:,0].numpy()+i*im_size+(i+1)*grid_border_size,
                    landmarks_batch[i,:,1].numpy()+grid_border_size,
                    s = 10,marker = '.',c = 'r')
        plt.title("Batch from dataloader")

for i_batch,sample_batch in enumerate(dataloader):
    print(i_batch,
          sample_batch['image'].size(),
          sample_batch['landmarks'].size())

    if i_batch == 3:
        plt.figure()
        show_landmarks_batch(sample_batch)
        plt.axis("off")
        break
plt.show()

out

*** person-7.jpg
Num, Image_Shape, Landmarks_Shape
0 (324, 215, 3) (68, 2)
1 (500, 333, 3) (68, 2)
2 (250, 258, 3) (68, 2)
3 (434, 290, 3) (68, 2)
0 torch.Size([3, 224, 224]) torch.Size([68, 2])
1 torch.Size([3, 224, 224]) torch.Size([68, 2])
2 torch.Size([3, 224, 224]) torch.Size([68, 2])
3 torch.Size([3, 224, 224]) torch.Size([68, 2])
0 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])
1 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])
2 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])
3 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])

Process finished with exit code 0

---

推荐阅读：https://github.com/deepinsight/insightface