在目标检测或者分割的时候,我们需要同时对图像和对应的方框或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()