Source code for beacon_aug.properties

# Copyright 2021 Adobe
# All Rights Reserved.

# NOTICE: Adobe permits you to use, modify, and distribute this file in
# accordance with the terms of the Adobe license agreement accompanying
# it.


import cv2
import matplotlib.pyplot as plt
import matplotlib
from skimage import data
import os
import random
import time
import beacon_aug as BA
import torch
import numpy as np
from typing import Callable, Iterable
import inspect
import cv2
import torchvision.transforms.functional as torch_f

kw_size_list = ["size", "crop_size", "kernel_size"]
kw_wh_list = ["weight", "heigh"]


[docs]def set_seed(seed): ''' Set Random seed to control the replication of transformations ''' random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed)
[docs]def isOpDifferentiable(op=None, size: Iterable[int] = (2, 3, 8, 8), # rgb image (8x8) batch size 2, device: str = "cpu"): """ Check if `op` is differentiable. `source <https://github.com/gist/holdgeof/22a496b76758cd9bfb690c8e88ac1c92>`_ Contributed by Yannick Hold-Geoffroy e.g. .. code-block:: import beacon_aug as BA BA.properties.isOpDifferentiable(op=BA.RandomCrop(library="torch")) """ # set default if op == None: op = BA.RandomCrop(library="torch") x = torch.rand(size, requires_grad=True, device=device) try: # make sure the size is smaller than the input image for kw in kw_size_list: if kw in dir(op.processor): setattr(op.processor, kw, (3, 3)) for kw in kw_wh_list: if kw in dir(op.processor): setattr(op.processor, kw, 3) # set "as_layter = True" to support tensor input in Beacon_aug y = op.apply(x, as_layer=True) loss = y.abs().sum() loss.backward() except Exception as e: print(e) return False is_diff = ( x.grad.shape == tuple(size) and torch.any(y != 0.) and torch.all(torch.isfinite(x.grad)).item() ) return is_diff
[docs]def isAntiAliasing(BA_operator, library, interpolation, **kwargs): """ Check if `op` is anti aliasing. `source Paper <https://github.com/GaParmar/clean-fid>`_ Whether or not the operator is anti-alising depends on the PSNR between original and rescaled images Contributed by `Richard Zhang`, e.g. .. code-block:: import beacon_aug as BA BA.properties.isAntiAliasing(BA.Resize, library= "torchvision",interpolation = "linear" ) """ def psnr_uint8(img0, img1): return -10*np.log10(np.mean((img0/255. - img1/255.)**2)) is_anti_aliasing = False img_np = np.zeros((128, 128, 3), dtype='uint8')+255 img_np = cv2.circle(img_np, (128//2, 128//2), 50, (0, 0, 0), 1) if "resize" in BA_operator.__name__.lower(): # interpolation_all = ["nearest", "linear", "area", "cubic", "Lanczos", "hamming"] interpolation_dict = {"albumentations": {"nearest": cv2.INTER_NEAREST, "linear": cv2.INTER_LINEAR, "area": cv2.INTER_AREA, "cubic": cv2.INTER_CUBIC, "Lanczos": cv2.INTER_LANCZOS4, }, "imgaug": {"nearest": "nearest", "linear": "linear", "cubic": "cubic", "area": "area" }, "torchvision": {"nearest": torch_f.InterpolationMode.NEAREST, "linear": torch_f.InterpolationMode.BILINEAR, "box": torch_f.InterpolationMode.BOX, "cubic": torch_f.InterpolationMode.BICUBIC, "Lanczos": torch_f.InterpolationMode.LANCZOS, "hamming": torch_f.InterpolationMode.HAMMING, }} interp = interpolation_dict[library][interpolation] if library in interpolation_dict else interpolation # print(interp) op = BA.Resize(p=1, library=library, interpolation=interp, height=16, width=16, **kwargs) op_up = BA.Resize(p=1, library=library, interpolation=interp, height=128, width=128) img_resized = op(image=img_np)["image"].copy() # downscale to 16x16 img_up = op_up(image=img_resized)["image"].copy() # upscale to 128*128 psnr = psnr_uint8(img_np, img_up) if psnr > 18: is_anti_aliasing = True return is_anti_aliasing
[docs]def library_attributes(BA_operator, image_data=None, **kwargs): ''' Visualize the augmentation result comparision to all available libraries e.g. .. code-block:: import beacon_aug as BA attributes_result = library_attributes(BA.Add) ''' if image_data == None: image_data = np.random.randint(255, size=(8, 8, 3)) avail_libraries = BA_operator(**kwargs).avail_libraries attributes_result = {"runtime": {}, "differentiable": {}} for i, library in enumerate(avail_libraries): t_before = time.time() op = BA_operator(always_apply=False, p=1, library=library, **kwargs) image_auged = op(image=image_data)["image"] t_after = time.time() runtime = t_after - t_before attributes_result["runtime"][library] = runtime attributes_result["differentiable"][library] = isOpDifferentiable(op) return attributes_result