mirror of
https://github.com/facebookresearch/moco-v3.git
synced 2025-06-03 14:59:22 +08:00
79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
# All rights reserved.
|
|
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
from PIL import Image, ImageFilter, ImageOps
|
|
import math
|
|
import random
|
|
import torchvision.transforms.functional as tf
|
|
|
|
|
|
class TwoCropsTransform:
|
|
"""Take two random crops of one image"""
|
|
|
|
def __init__(self, base_transform1, base_transform2):
|
|
self.base_transform1 = base_transform1
|
|
self.base_transform2 = base_transform2
|
|
|
|
def __call__(self, x):
|
|
im1 = self.base_transform1(x)
|
|
im2 = self.base_transform2(x)
|
|
return [im1, im2]
|
|
|
|
|
|
class RandomResizedCropBYOL(object):
|
|
"""Random crop augmentation from BYOL: https://arxiv.org/abs/2006.07733"""
|
|
|
|
def __init__(self, size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.), interpolation=Image.BICUBIC):
|
|
self.size = (size, size)
|
|
if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
|
|
warnings.warn("range should be of kind (min, max)")
|
|
|
|
self.interpolation = interpolation
|
|
self.scale = scale
|
|
self.ratio = ratio
|
|
|
|
@staticmethod
|
|
def get_params(img, scale, ratio):
|
|
width, height = img.size
|
|
area = height * width
|
|
|
|
target_area = random.uniform(*scale) * area
|
|
log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
|
|
aspect_ratio = math.exp(random.uniform(*log_ratio))
|
|
|
|
w = int(round(math.sqrt(target_area * aspect_ratio)))
|
|
h = int(round(math.sqrt(target_area / aspect_ratio)))
|
|
|
|
w = min(w, width)
|
|
h = min(h, height)
|
|
|
|
i = random.randint(0, height - h)
|
|
j = random.randint(0, width - w)
|
|
|
|
return i, j, h, w
|
|
|
|
def __call__(self, img):
|
|
i, j, h, w = self.get_params(img, self.scale, self.ratio)
|
|
return tf.resized_crop(img, i, j, h, w, self.size, self.interpolation)
|
|
|
|
|
|
class GaussianBlur(object):
|
|
"""Gaussian blur augmentation from SimCLR: https://arxiv.org/abs/2002.05709"""
|
|
|
|
def __init__(self, sigma=[.1, 2.]):
|
|
self.sigma = sigma
|
|
|
|
def __call__(self, x):
|
|
sigma = random.uniform(self.sigma[0], self.sigma[1])
|
|
x = x.filter(ImageFilter.GaussianBlur(radius=sigma))
|
|
return x
|
|
|
|
|
|
class Solarize(object):
|
|
"""Solarize augmentation from BYOL: https://arxiv.org/abs/2006.07733"""
|
|
|
|
def __call__(self, x):
|
|
return ImageOps.solarize(x) |