"""Utilities for augmenting image data."""
import detectron2.data.transforms as T
import imgaug.augmenters as iaa
import numpy as np
import deepdisc.astrodet.detectron as detectron_addons
import random
[docs]
LAMBDA_EFFS = [3671,4827,6223,7546,8691,9712]
[docs]
A_EBV = np.array([4.81,3.64,2.70,2.06,1.58,1.31])
[docs]
def redden(image, rng_seed=None):
"""
Parameters
----------
image: ndarray
rng_seed : np.random.Generator
Random state that is seeded. if none, use machine entropy.
Returns
-------
augmented image
"""
new_ebv = np.random.uniform(0, 0.5)
image = np.float32(image*(10.**(-A_EBV*new_ebv/2.5)))
return image
[docs]
def gaussblur(image, rng_seed=None):
"""
Parameters
----------
image: ndarray
rng_seed : np.random.Generator
Random state that is seeded. if none, use machine entropy.
Returns
-------
augmented image
"""
if rng_seed is None:
rng_seed = np.random.default_rng()
aug = iaa.GaussianBlur(sigma=10, seed=rng_seed)
return aug.augment_image(image)
[docs]
def scale_psf(sigi, lambda_eff):
i_eff = 7546
sig_lambda = sigi* (lambda_eff)**(-0.3)/(i_eff**(-0.3))
return sig_lambda
[docs]
def multiband_gaussblur(image, rng_seed=None):
"""
Parameters
----------
image: ndarray
rng_seed : np.random.Generator
Random state that is seeded. if none, use machine entropy.
Returns
-------
augmented image
"""
if rng_seed is None:
rng_seed = np.random.default_rng()
imgs = np.zeros(image.shape, dtype=np.float32)
sigmai = random.random()
for i in range(6):
sigma = scale_psf(sigmai,LAMBDA_EFFS[i])
aug = iaa.GaussianBlur(sigma=sigma, seed=rng_seed)
imaug = aug.augment_image(image[:,:,i])
imgs[:,:,i] = imaug
return imgs
[docs]
def addelementwise16(image, rng_seed=None):
"""
Parameters
----------
image: ndarray
rng_seed : np.random.Generator
Random state that is seeded. if none, use machine entropy.
Returns
-------
augmented image
"""
if rng_seed is None:
rng_seed = np.random.default_rng()
aug = iaa.AddElementwise((-3276, 3276), seed=rng_seed)
return aug.augment_image(image)
[docs]
def addelementwise8(image, rng_seed=None):
"""
Parameters
----------
image: ndarray
rng_seed : np.random.Generator
Random state that is seeded. if none, use machine entropy.
Returns
-------
augmented image
"""
if rng_seed is None:
rng_seed = np.random.default_rng()
aug = iaa.AddElementwise((-25, 25), seed=rng_seed)
return aug.augment_image(image)
[docs]
def addelementwise(image, rng_seed=None):
"""
Parameters
----------
image: ndarray
rng_seed : np.random.Generator
Random state that is seeded. if none, use machine entropy.
Returns
-------
augmented image
"""
if rng_seed is None:
rng_seed = np.random.default_rng()
aug = iaa.AddElementwise((-image.max() * 0.1, image.max() * 0.1), seed=rng_seed)
return aug.augment_image(image)
[docs]
def centercrop(image):
"""Crop an image to just the center portion
Parameters
----------
image: ndarray
Returns
-------
cropped image
"""
h, w = image.shape[:2]
hc = (h - h // 2) // 2
wc = (w - w // 2) // 2
image = image[hc : hc + h // 2, wc : wc + w // 2]
return image
[docs]
def train_augs(image):
"""Get the augmentation list
Parameters
----------
image: image
The image to be augmented
Returns
-------
augs: detectron_addons.KRandomAugmentationList
The list of augs for training. Set to RandomRotation, RandomFlip, RandomCrop
"""
augs = detectron_addons.KRandomAugmentationList(
[
# my custom augs
T.RandomRotation([-90, 90, 180], sample_style="choice"),
T.RandomFlip(prob=0.5),
T.RandomFlip(prob=0.5, horizontal=False, vertical=True),
],
k=-1,
cropaug=T.RandomCrop("relative", (0.5, 0.5)),
)
return augs
[docs]
def dc2_train_augs(image):
"""Get the augmentation list
Parameters
----------
image: image
The image to be augmented
Returns
-------
augs: detectron_addons.KRandomAugmentationList
The list of augs for training. Set to RandomRotation, RandomFlip, RandomCrop
"""
augs = detectron_addons.KRandomAugmentationList(
[
# my custom augs
T.RandomRotation([-90, 90, 180], sample_style="choice"),
T.RandomFlip(prob=0.5),
T.RandomFlip(prob=0.5, horizontal=False, vertical=True),
#detectron_addons.CustomAug(multiband_gaussblur,prob=1.0),
#detectron_addons.CustomAug(redden,prob=1.0),
],
k=-1,
cropaug=None,
#cropaug=T.RandomCrop("relative", (0.5, 0.5))
)
return augs
[docs]
def dc2_train_augs_full(image):
"""Get the augmentation list
Parameters
----------
image: image
The image to be augmented
Returns
-------
augs: detectron_addons.KRandomAugmentationList
The list of augs for training. Set to RandomRotation, RandomFlip, RandomCrop
"""
augs = detectron_addons.KRandomAugmentationList(
[
# my custom augs
T.RandomRotation([-90, 90, 180], sample_style="choice"),
T.RandomFlip(prob=0.5),
T.RandomFlip(prob=0.5, horizontal=False, vertical=True),
detectron_addons.CustomAug(multiband_gaussblur,prob=1.0),
detectron_addons.CustomAug(redden,prob=1.0),
],
k=-1,
cropaug=None,
#cropaug=T.RandomCrop("relative", (0.5, 0.5))
)
return augs
[docs]
def hsc_test_augs(image):
"""Get the augmentation list
Parameters
----------
image: image
The image to be augmented
Returns
-------
augs: detectron2 AugmentationList
The augs for hsc testing. Set to 50% Crop due to memory constraints
"""
augs = T.AugmentationList(
[
T.CropTransform(
image.shape[1] // 4,
image.shape[0] // 4,
image.shape[1] // 2,
image.shape[0] // 2,
)
]
)
return augs