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()

在这里插入图片描述

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章