Pytorch 圖片載入

目錄

• 載入圖片和坐標
• Transforms
• Compose transforms
• 數據集的迭代

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 Dataset, DataLoader
from torchvision import transforms, utils

# Ignore warnings
import warnings
warnings.filterwarnings("ignore")

載入圖片和坐標

landmarks_frame = pd.read_csv('data/faces/face_landmarks.csv')
n=65
img_name = landmarks_frame.iloc[n, 0]  #獲取圖片的名稱
landmarks = landmarks_frame.iloc[n, 1:].as_matrix()  #獲取點的位置
landmarks = landmarks.astype('float').reshape(-1, 2)

landmarks_frame.iloc[:3, :] #展示一下csv里面的格式

	image_name	part_0_x	part_0_y	part_1_x	part_1_y	part_2_x	part_2_y	part_3_x	part_3_y	part_4_x	...	part_63_x	part_63_y	part_64_x	part_64_y	part_65_x	part_65_y	part_66_x	part_66_y	part_67_x	part_67_y
0	0805personali01.jpg	27	83	27	98	29	113	33	127	39	...	93	136	100	141	93	135	89	135	84	134
1	1084239450_e76e00b7e7.jpg	70	236	71	257	75	278	82	299	90	...	148	311	179	308	149	312	137	314	128	312
2	10comm-decarlo.jpg	66	114	65	128	67	142	68	156	72	...	128	162	136	167	127	166	121	165	116	164

3 rows × 137 columns

接下來，是如何展示圖片，以及把點畫在圖片之上

def show_landmarks(image, landmarks):
    fig, ax = plt.subplots()
    ax.imshow(image)
    ax.scatter(landmarks[:, 0], landmarks[:, 1], s=10, marker='.', c='r')
    plt.pause(0.001)  #暫停讓圖片更新?
    plt.show()
    
show_landmarks(io.imread(os.path.join('data/faces/', img_name)),
              landmarks)
    

torch.utils.data.Dataset是一個抽象基類表示一個數據集，我們需要為其設定__len__方法和__getitem__方法.

class FaceLandmarksDataset(Dataset):
    def __init__(self, csv_file, root_dir, transform=None):
        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

利用這個類，我們來展示一下前4幅圖像

face_dataset = FaceLandmarksDataset(csv_file='data/faces/face_landmarks.csv', 
                                    root_dir='data/faces/')

def show_landmarks(image, landmarks):
    plt.imshow(image)
    plt.scatter(landmarks[:, 0], landmarks[:, 1], s=10, marker='.', c='r')
    plt.pause(0.001)  #暫停讓圖片更新?
    
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

0 (324, 215, 3) (68, 2)

1 (500, 333, 3) (68, 2)

2 (250, 258, 3) (68, 2)

3 (434, 290, 3) (68, 2)

Transforms

很多時候，我們需要對圖片進行一些變化，比方說大小的調整等等

利用函子(_call_)能夠很好很方便對圖片進行處理

class Rescale(object):
    """Rescale the image in a sample to a given size.

    Args:
        output_size (tuple or int): Desired output size. If tuple, output is
            matched to output_size. If int, smaller of image edges is matched
            to output_size keeping aspect ratio the same.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple)) #output_size應當是一個整數或者元組
        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
            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))

        # 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] #坐標也要相應改變大小

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

class RandomCrop(object): #隨機裁剪，但是實際上是一整塊來的
    """Crop randomly the image in a sample.

    Args:
        output_size (tuple or int): Desired output size. If int, square crop
            is made.
    """

    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

    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}

class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

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

        # swap color axis because
        # numpy image: H x W x C
        # torch image: C X H X W
        image = image.transpose((2, 0, 1)) #把ndarray轉換為tensor需要改變順序
        return {'image': torch.from_numpy(image),
                'landmarks': torch.from_numpy(landmarks)}

Compose transforms

利用torchvision.transforms.Compose可以幫助我們對一個圖片進行多個操作

scale = Rescale(256)
crop = RandomCrop(128)
composed = transforms.Compose([Rescale(256),
                               RandomCrop(224)])

# Apply each of the above transforms on sample.
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()

數據集的迭代

我們可以用 for ... in ... 來迭代數據集，但是這么做並不方便，因為很多時候訓練神經網絡是要分批和打亂順序的torch.utils.data.DataLoader可以幫助我們完成這一個目標

transformed_dataset = FaceLandmarksDataset(csv_file='data/faces/face_landmarks.csv',
                                           root_dir='data/faces/',
                                           transform=transforms.Compose([
                                               Rescale(256),
                                               RandomCrop(224),
                                               ToTensor()
                                           ]))
dataloader = DataLoader(transformed_dataset, batch_size=4,
                        shuffle=True, num_workers=0) #batch_size: batch的大小 shuffle=True表示順序打亂

def show_landmarks_batch(sample_batched):
    """Show image with landmarks for a batch of samples."""
    images_batch, landmarks_batch = \
            sample_batched['image'], sample_batched['landmarks']
    batch_size = len(images_batch)
    im_size = images_batch.size(2)
    grid_border_size = 2

    grid = utils.make_grid(images_batch) #為圖片加入邊框
    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_batched in enumerate(dataloader):
    print(i_batch, sample_batched['image'].size(),
          sample_batched['landmarks'].size())


    if i_batch == 0:
        plt.figure()
        show_landmarks_batch(sample_batched)
        plt.axis('off')
        plt.ioff()
        plt.show()
        break

0 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])
224