pytorch中如何同時對image，bounding box, instance mask 進行相同的圖像變換？

在目標檢測或者分割的時候，我們需要同時對圖像和對應的方框或mask進行相同的變換，然後作爲ground truth訓練模型。

pytorch提供這樣的服務，且十分簡單，只要自己定義自己的函數即可，然後調用。
這裏以旋轉爲例：

def my_transform(image, mask):
    
    image = F.to_pil_image(image)
    mask = F.to_pil_image(mask)
    if random.random() > 0.1:
#         angle = random.randint(-20, 20)
        angle = 20
        image = F.rotate(image, angle)
        mask = F.rotate(mask, angle)
    return image, mask

其他的變換類似於這樣，只要稍加變換函數裏面的功能即可。

下面進行展示所有的實例代碼和結果

import random
import numpy as np
import torch
import torchvision.transforms.functional as F
import matplotlib.pyplot as plt
%matplotlib inline

def my_transform(image, mask):
    
    image = F.to_pil_image(image)
    mask = F.to_pil_image(mask)
    if random.random() > 0.1:
#         angle = random.randint(-20, 20)
        angle = 20
        image = F.rotate(image, angle)
        mask = F.rotate(mask, angle)
    return image, mask


name = 'image/000000468124.jpg'
mask_name = 'image/000000468124_mask_0.jpg'
img = cv2.imread(name)
mask = cv2.imread(mask_name)

plt.figure()
plt.subplot(221)
plt.imshow(img)
plt.axis('off')
plt.subplot(222)
plt.imshow(mask, cmap='gray')
plt.axis('off')

img, mask = my_transform(img, mask)
plt.subplot(223)
plt.imshow(img)
plt.axis('off')
plt.subplot(224)
plt.imshow(mask, cmap='gray')
plt.axis('off')

plt.show()