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

在这里插入图片描述

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值