def test_compose_with_keypoint_noop(keypoints, keypoint_format, labels):
image = np.ones((100, 100, 3))
if labels is not None:
aug = Compose([NoOp(p=1.0)],
keypoint_params={
"format": keypoint_format,
"label_fields": ["labels"]
})
transformed = aug(image=image, keypoints=keypoints, labels=labels)
else:
aug = Compose([NoOp(p=1.0)],
keypoint_params={"format": keypoint_format})
transformed = aug(image=image, keypoints=keypoints)
assert np.array_equal(transformed["image"], image)
assert transformed["keypoints"] == keypoints
class ImgAugTTATransform(AugmentationFactoryBase):
TRANSFORM_LIST = (
NoOp(p=1),
HorizontalFlip(p=1),
VerticalFlip(p=1),
Rotate(limit=(180, 180), p=1),
# Lambda(RollingY(shift=(-20, 20)).transform),
# Lambda(RollingY(shift=(-20, 20)).transform),
)
CURR_TRANSFORM = TRANSFORM_LIST[0]
def __init__(self, image_size=(126, 32)):
self.image_size = image_size
def build_train(self):
pass
def build_test(self):
print("Using {}".format(self.CURR_TRANSFORM))
return Compose([
Lambda(minmax_norm),
self.CURR_TRANSFORM,
Resize(
self.image_size[0],
self.image_size[1],
interpolation=cv2.INTER_CUBIC,
),
Lambda(normalize),
ToTensor(),
])
def test_compose_with_keypoint_noop_error_label_fields(keypoints,
keypoint_format):
image = np.ones((100, 100, 3))
aug = Compose([NoOp(p=1.0)],
keypoint_params={
"format": keypoint_format,
"label_fields": "class_id"
})
with pytest.raises(Exception):
aug(image=image, keypoints=keypoints, cls_id=[0])
def test_compose_with_keypoint_noop_label_outside(keypoints, keypoint_format,
labels):
image = np.ones((100, 100, 3))
aug = Compose([NoOp(p=1.0)],
keypoint_params={
"format": keypoint_format,
"label_fields": list(labels.keys())
})
transformed = aug(image=image, keypoints=keypoints, **labels)
assert np.array_equal(transformed["image"], image)
assert transformed["keypoints"] == keypoints
for k, v in labels.items():
assert transformed[k] == v
def get_dls(train_img_paths,
train_mask_paths,
val_img_paths,
val_mask_paths,
size=128,
batch_size=1,
num_workers=1,
pin_memory=False):
trn_img, trn_mask = create_arrays(train_img_paths, train_mask_paths, size)
val_img, val_mask = create_arrays(val_img_paths, val_mask_paths, size)
trn_ds = SaltDatasetArrays(trn_img, trn_mask, tfms=dihedral)
val_ds = SaltDatasetArrays(val_img, val_mask, tfms=NoOp())
trn_dl = DataLoader(trn_ds,
num_workers=num_workers,
batch_size=batch_size,
pin_memory=pin_memory)
val_dl = DataLoader(val_ds,
num_workers=num_workers,
batch_size=batch_size,
pin_memory=pin_memory)
return trn_dl, val_dl
def transform(image, mask, image_name, mask_name):
x, y = image, mask
rand = random.uniform(0, 1)
if (rand > 0.5):
images_name = [f"{image_name}"]
masks_name = [f"{mask_name}"]
images_aug = [x]
masks_aug = [y]
it = iter(images_name)
it2 = iter(images_aug)
imagedict = dict(zip(it, it2))
it = iter(masks_name)
it2 = iter(masks_aug)
masksdict = dict(zip(it, it2))
return imagedict, masksdict
mask_density = np.count_nonzero(y)
## Augmenting only images with Gloms
if (mask_density > 0):
try:
h, w, c = x.shape
except Exception as e:
image = image[:-1]
x, y = image, mask
h, w, c = x.shape
aug = Blur(p=1, blur_limit=3)
augmented = aug(image=x, mask=y)
x0 = augmented['image']
y0 = augmented['mask']
# aug = CenterCrop(p=1, height=32, width=32)
# augmented = aug(image=x, mask=y)
# x1 = augmented['image']
# y1 = augmented['mask']
## Horizontal Flip
aug = HorizontalFlip(p=1)
augmented = aug(image=x, mask=y)
x2 = augmented['image']
y2 = augmented['mask']
aug = VerticalFlip(p=1)
augmented = aug(image=x, mask=y)
x3 = augmented['image']
y3 = augmented['mask']
# aug = Normalize(p=1)
# augmented = aug(image=x, mask=y)
# x4 = augmented['image']
# y4 = augmented['mask']
aug = Transpose(p=1)
augmented = aug(image=x, mask=y)
x5 = augmented['image']
y5 = augmented['mask']
aug = RandomGamma(p=1)
augmented = aug(image=x, mask=y)
x6 = augmented['image']
y6 = augmented['mask']
## Optical Distortion
aug = OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5)
augmented = aug(image=x, mask=y)
x7 = augmented['image']
y7 = augmented['mask']
## Grid Distortion
aug = GridDistortion(p=1)
augmented = aug(image=x, mask=y)
x8 = augmented['image']
y8 = augmented['mask']
aug = RandomGridShuffle(p=1)
augmented = aug(image=x, mask=y)
x9 = augmented['image']
y9 = augmented['mask']
aug = HueSaturationValue(p=1)
augmented = aug(image=x, mask=y)
x10 = augmented['image']
y10 = augmented['mask']
# aug = PadIfNeeded(p=1)
# augmented = aug(image=x, mask=y)
# x11 = augmented['image']
# y11 = augmented['mask']
aug = RGBShift(p=1)
augmented = aug(image=x, mask=y)
x12 = augmented['image']
y12 = augmented['mask']
## Random Brightness
aug = RandomBrightness(p=1)
augmented = aug(image=x, mask=y)
x13 = augmented['image']
y13 = augmented['mask']
## Random Contrast
aug = RandomContrast(p=1)
augmented = aug(image=x, mask=y)
x14 = augmented['image']
y14 = augmented['mask']
#aug = MotionBlur(p=1)
#augmented = aug(image=x, mask=y)
# x15 = augmented['image']
# y15 = augmented['mask']
aug = MedianBlur(p=1, blur_limit=5)
augmented = aug(image=x, mask=y)
x16 = augmented['image']
y16 = augmented['mask']
aug = GaussianBlur(p=1, blur_limit=3)
augmented = aug(image=x, mask=y)
x17 = augmented['image']
y17 = augmented['mask']
aug = GaussNoise(p=1)
augmented = aug(image=x, mask=y)
x18 = augmented['image']
y18 = augmented['mask']
aug = GlassBlur(p=1)
augmented = aug(image=x, mask=y)
x19 = augmented['image']
y19 = augmented['mask']
aug = CLAHE(clip_limit=1.0,
tile_grid_size=(8, 8),
always_apply=False,
p=1)
augmented = aug(image=x, mask=y)
x20 = augmented['image']
y20 = augmented['mask']
aug = ChannelShuffle(p=1)
augmented = aug(image=x, mask=y)
x21 = augmented['image']
y21 = augmented['mask']
aug = ToGray(p=1)
augmented = aug(image=x, mask=y)
x22 = augmented['image']
y22 = augmented['mask']
aug = ToSepia(p=1)
augmented = aug(image=x, mask=y)
x23 = augmented['image']
y23 = augmented['mask']
aug = JpegCompression(p=1)
augmented = aug(image=x, mask=y)
x24 = augmented['image']
y24 = augmented['mask']
aug = ImageCompression(p=1)
augmented = aug(image=x, mask=y)
x25 = augmented['image']
y25 = augmented['mask']
aug = Cutout(p=1)
augmented = aug(image=x, mask=y)
x26 = augmented['image']
y26 = augmented['mask']
# aug = CoarseDropout(p=1, max_holes=8, max_height=32, max_width=32)
# augmented = aug(image=x, mask=y)
# x27 = augmented['image']
# y27 = augmented['mask']
# aug = ToFloat(p=1)
# augmented = aug(image=x, mask=y)
# x28 = augmented['image']
# y28 = augmented['mask']
aug = FromFloat(p=1)
augmented = aug(image=x, mask=y)
x29 = augmented['image']
y29 = augmented['mask']
## Random Brightness and Contrast
aug = RandomBrightnessContrast(p=1)
augmented = aug(image=x, mask=y)
x30 = augmented['image']
y30 = augmented['mask']
aug = RandomSnow(p=1)
augmented = aug(image=x, mask=y)
x31 = augmented['image']
y31 = augmented['mask']
aug = RandomRain(p=1)
augmented = aug(image=x, mask=y)
x32 = augmented['image']
y32 = augmented['mask']
aug = RandomFog(p=1)
augmented = aug(image=x, mask=y)
x33 = augmented['image']
y33 = augmented['mask']
aug = RandomSunFlare(p=1)
augmented = aug(image=x, mask=y)
x34 = augmented['image']
y34 = augmented['mask']
aug = RandomShadow(p=1)
augmented = aug(image=x, mask=y)
x35 = augmented['image']
y35 = augmented['mask']
aug = Lambda(p=1)
augmented = aug(image=x, mask=y)
x36 = augmented['image']
y36 = augmented['mask']
aug = ChannelDropout(p=1)
augmented = aug(image=x, mask=y)
x37 = augmented['image']
y37 = augmented['mask']
aug = ISONoise(p=1)
augmented = aug(image=x, mask=y)
x38 = augmented['image']
y38 = augmented['mask']
aug = Solarize(p=1)
augmented = aug(image=x, mask=y)
x39 = augmented['image']
y39 = augmented['mask']
aug = Equalize(p=1)
augmented = aug(image=x, mask=y)
x40 = augmented['image']
y40 = augmented['mask']
aug = Posterize(p=1)
augmented = aug(image=x, mask=y)
x41 = augmented['image']
y41 = augmented['mask']
aug = Downscale(p=1)
augmented = aug(image=x, mask=y)
x42 = augmented['image']
y42 = augmented['mask']
aug = MultiplicativeNoise(p=1)
augmented = aug(image=x, mask=y)
x43 = augmented['image']
y43 = augmented['mask']
aug = FancyPCA(p=1)
augmented = aug(image=x, mask=y)
x44 = augmented['image']
y44 = augmented['mask']
# aug = MaskDropout(p=1)
# augmented = aug(image=x, mask=y)
# x45 = augmented['image']
# y45 = augmented['mask']
aug = GridDropout(p=1)
augmented = aug(image=x, mask=y)
x46 = augmented['image']
y46 = augmented['mask']
aug = ColorJitter(p=1)
augmented = aug(image=x, mask=y)
x47 = augmented['image']
y47 = augmented['mask']
## ElasticTransform
aug = ElasticTransform(p=1,
alpha=120,
sigma=512 * 0.05,
alpha_affine=512 * 0.03)
augmented = aug(image=x, mask=y)
x50 = augmented['image']
y50 = augmented['mask']
aug = CropNonEmptyMaskIfExists(p=1, height=22, width=32)
augmented = aug(image=x, mask=y)
x51 = augmented['image']
y51 = augmented['mask']
aug = IAAAffine(p=1)
augmented = aug(image=x, mask=y)
x52 = augmented['image']
y52 = augmented['mask']
# aug = IAACropAndPad(p=1)
# augmented = aug(image=x, mask=y)
# x53 = augmented['image']
# y53 = augmented['mask']
aug = IAAFliplr(p=1)
augmented = aug(image=x, mask=y)
x54 = augmented['image']
y54 = augmented['mask']
aug = IAAFlipud(p=1)
augmented = aug(image=x, mask=y)
x55 = augmented['image']
y55 = augmented['mask']
aug = IAAPerspective(p=1)
augmented = aug(image=x, mask=y)
x56 = augmented['image']
y56 = augmented['mask']
aug = IAAPiecewiseAffine(p=1)
augmented = aug(image=x, mask=y)
x57 = augmented['image']
y57 = augmented['mask']
aug = LongestMaxSize(p=1)
augmented = aug(image=x, mask=y)
x58 = augmented['image']
y58 = augmented['mask']
aug = NoOp(p=1)
augmented = aug(image=x, mask=y)
x59 = augmented['image']
y59 = augmented['mask']
# aug = RandomCrop(p=1, height=22, width=22)
# augmented = aug(image=x, mask=y)
# x61 = augmented['image']
# y61 = augmented['mask']
# aug = RandomResizedCrop(p=1, height=22, width=20)
# augmented = aug(image=x, mask=y)
# x63 = augmented['image']
# y63 = augmented['mask']
aug = RandomScale(p=1)
augmented = aug(image=x, mask=y)
x64 = augmented['image']
y64 = augmented['mask']
# aug = RandomSizedCrop(p=1, height=22, width=20, min_max_height = [32,32])
# augmented = aug(image=x, mask=y)
# x66 = augmented['image']
# y66 = augmented['mask']
# aug = Resize(p=1, height=22, width=20)
# augmented = aug(image=x, mask=y)
# x67 = augmented['image']
# y67 = augmented['mask']
aug = Rotate(p=1)
augmented = aug(image=x, mask=y)
x68 = augmented['image']
y68 = augmented['mask']
aug = ShiftScaleRotate(p=1)
augmented = aug(image=x, mask=y)
x69 = augmented['image']
y69 = augmented['mask']
aug = SmallestMaxSize(p=1)
augmented = aug(image=x, mask=y)
x70 = augmented['image']
y70 = augmented['mask']
images_aug.extend([
x, x0, x2, x3, x5, x6, x7, x8, x9, x10, x12, x13, x14, x16, x17,
x18, x19, x20, x21, x22, x23, x24, x25, x26, x29, x30, x31, x32,
x33, x34, x35, x36, x37, x38, x39, x40, x41, x42, x43, x44, x46,
x47, x50, x51, x52, x54, x55, x56, x57, x58, x59, x64, x68, x69,
x70
])
masks_aug.extend([
y, y0, y2, y3, y5, y6, y7, y8, y9, y10, y12, y13, y14, y16, y17,
y18, y19, y20, y21, y22, y23, y24, y25, y26, y29, y30, y31, y32,
y33, y34, y35, y36, y37, y38, y39, y40, y41, y42, y43, y44, y46,
y47, y50, y51, y52, y54, y55, y56, y57, y58, y59, y64, y68, y69,
y70
])
idx = -1
images_name = []
masks_name = []
for i, m in zip(images_aug, masks_aug):
if idx == -1:
tmp_image_name = f"{image_name}"
tmp_mask_name = f"{mask_name}"
else:
tmp_image_name = f"{image_name}_{smalllist[idx]}"
tmp_mask_name = f"{mask_name}_{smalllist[idx]}"
images_name.extend(tmp_image_name)
masks_name.extend(tmp_mask_name)
idx += 1
it = iter(images_name)
it2 = iter(images_aug)
imagedict = dict(zip(it, it2))
it = iter(masks_name)
it2 = iter(masks_aug)
masksdict = dict(zip(it, it2))
return imagedict, masksdict
import torch
from pathlib import Path
from tqdm import trange
from albumentations import Resize, NoOp
from iharm.data.hdataset import HDataset
from iharm.data.transforms import HCompose, LongestMaxSizeIfLarger
from iharm.inference.predictor import Predictor
from iharm.inference.metrics import MetricsHub, MSE, fMSE, PSNR, N
from iharm.inference.utils import load_model, find_checkpoint
from iharm.mconfigs import ALL_MCONFIGS
from iharm.utils.exp import load_config_file
from iharm.utils.log import logger, add_new_file_output_to_logger
RESIZE_STRATEGIES = {
'None': NoOp(),
'LimitLongest1024': LongestMaxSizeIfLarger(1024),
'Fixed256': Resize(256, 256),
'Fixed512': Resize(512, 512)
}
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('model_type', choices=ALL_MCONFIGS.keys())
parser.add_argument(
'checkpoint',
type=str,
help='The path to the checkpoint. '
'This can be a relative path (relative to cfg.MODELS_PATH) '
def create_arrays(img_paths, mask_paths, size=None):
tfm = NoOp() if size is None else Resize(size, size)
tfm = Compose([tfm, Normalize()])
imgs = np.stack([open_image(i, tfm=tfm) for i in img_paths])
masks = np.stack([open_mask(i, tfm=tfm) for i in mask_paths])
return imgs, masks
def open_mask(mask_path, tfm=NoOp()):
mask = cv2.cvtColor(cv2.imread(mask_path), cv2.COLOR_BGR2RGB)
mask = (mask == 255).astype(np.float32)
mask = tfm(image=mask, mask=mask)['mask']
return np.prod(mask, axis=-1)
def open_image(img_path, tfm=NoOp()):
img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
img = img.astype(np.float32)
img = tfm(image=img)['image']
return img
def _FastRandAugment2d(img_size, set='general'):
_sets = {
'general': [
OneOf([
Blur(p=1),
GaussNoise(p=1),
MotionBlur(p=1),
MedianBlur(p=1),
]),
HorizontalFlip(p=1),
Compose([
PadIfNeeded(min_height=img_size * 2,
min_width=img_size * 2,
border_mode=cv2.BORDER_WRAP,
p=1),
RandomCrop(img_size, img_size, p=1),
],
p=1),
RandomBrightnessContrast(p=1),
InvertImg(p=1),
ToGray(p=1),
RGBShift(p=1),
OneOf([
GridDropout(holes_number_x=4,
holes_number_y=4,
random_offset=True,
ratio=r,
fill_value=i,
p=1) for i in _color_fill
for r in np.linspace(0.3, 0.7, 4)
],
p=1),
NoOp(p=1),
Equalize(p=1),
Rotate(45, border_mode=cv2.BORDER_WRAP, p=1),
ShiftScaleRotate(shift_limit=0.25,
scale_limit=0.2,
rotate_limit=45,
border_mode=cv2.BORDER_WRAP,
p=1),
Solarize(p=1),
Posterize(p=1),
IAAPerspective(p=1),
HueSaturationValue(hue_shift_limit=20,
sat_shift_limit=50,
val_shift_limit=50,
p=1),
RandomGamma(p=1),
ChannelShuffle(p=1),
OneOf([
CoarseDropout(max_holes=1,
max_height=img_size // 2,
max_width=img_size // 2,
fill_value=i,
p=1) for i in _color_fill
],
p=1),
OneOf([
CLAHE(clip_limit=2),
IAASharpen(),
IAAEmboss(),
], p=1),
],
}
return _sets[set]