pytorch transforms库如何自定义crop图像

【分享贴】基于 transforms库 自定义crop图像

pytorch transforms 库的默认用法

from torchvision import transforms
from PIL import Image
transform = transforms.Compose([
        transforms.Resize((680,680)),
        transforms.CenterCrop((256,256)),
        transforms.ToTensor(),
        transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])
img = Image.open("./test.jpg").convert("RGB")
img = transform(img) #[3*256*256]

通过以上代码就可以完成对输入进来的 PIL 格式图像的预处理

transforms 自定义crop

transforms 默认提供了 CenterCrop , RandomCrop, FiveCrop, TenCrop 多种裁剪方式，后三种可以满足在训练网络阶段完成图像的数据增广。

但是我们有没有想过如何自定义 crop 图像：

一、例如只裁剪图像的左上角，代码实现如下：

from torchvision.transforms.functional import crop
from torchvision import transforms
from PIL import Image

def crop_left_upper(image):
    return crop(image, 0, 0, 256, 256)
    
transform = transforms.Compose([
        transforms.Resize((680,680)),
        transforms.Lambda(crop_left_upper),
        transforms.ToTensor(),
        transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])
img = Image.open("./test.jpg").convert("RGB")
img = transform(img) #[3*256*256]

二、API 提供 FiveCrop，但我们也许需求在一张图中每次crop出2个tensor，可以这样写：

例如，我们根据图像规律发现，图像中左上角和中心区域是最具”信息“的地方，我们可以着重crop

from torchvision.transforms.functional import crop
from torchvision import transforms
from PIL import Image
transform = transforms.Compose([
        transforms.Resize((680,680)),
        transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop(crops, 0,0,256,256)), transforms.ToTensor()(crop(crops, 212,212,256,256))])),
        transforms.Lambda(lambda crops: torch.stack([transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])(crop) for crop in crops]))
])
img = Image.open("./test.jpg").convert("RGB")
img = transform(img) #[2*3*256*256]

当然，借助lambda匿名函数，还可以有更优的写法。

这样我们通过对 transforms 预处理模块的 API 小改造，以满足各式各样的图像预处理需求，由于该过程我们不使用cv2.，pillow等crop函数的参与，保证了torch代码的可读性与一致性