def generate_transforms2(img_size):
train_transform = Compose([
# A.RandomCrop(p=1, height=img_size, width=img_size),
A.Resize(height=img_size, width=img_size),
A.RandomSunFlare(p=1),
A.RandomFog(p=1),
A.RandomBrightness(p=1),
A.Rotate(p=1, limit=90),
A.RGBShift(p=1),
A.RandomSnow(p=1),
A.HorizontalFlip(p=1),
A.VerticalFlip(p=1),
A.RandomContrast(limit=0.5, p=1),
A.HueSaturationValue(p=1,
hue_shift_limit=20,
sat_shift_limit=30,
val_shift_limit=50),
# A.Cutout(p=1),
# A.Transpose(p=1),
A.JpegCompression(p=1),
A.CoarseDropout(p=1),
A.IAAAdditiveGaussianNoise(loc=0,
scale=(2.5500000000000003, 12.75),
per_channel=False,
p=1),
A.IAAAffine(scale=1.0,
translate_percent=None,
translate_px=None,
rotate=0.0,
shear=0.0,
order=1,
cval=0,
mode='reflect',
p=1),
A.IAAAffine(rotate=90., p=1),
A.IAAAffine(rotate=180., p=1),
Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
max_pixel_value=255.0,
p=1.0),
ToTensorV2()
])
val_transform = Compose([
Resize(height=img_size, width=img_size),
Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
max_pixel_value=255.0,
p=1.0),
ToTensorV2(),
])
return {"train": train_transform, "val": val_transform}
def get_transforms(
input_size, target_size, aug_scale_limit, aug_rot_limit, aug_shear_limit,
aug_brightness_limit, aug_contrast_limit):
crop_size = int(input_size * 0.9)
base_size = int(target_size * 1.14)
train_transforms = albm.Compose([
albm.IAAAffine(
scale=aug_scale_limit,
rotate=aug_rot_limit,
shear=aug_shear_limit,
p=1.0),
albm.CenterCrop(crop_size, crop_size),
albm.Resize(base_size, base_size),
albm.RandomCrop(target_size, target_size),
albm.RandomBrightnessContrast(
brightness_limit=aug_brightness_limit,
contrast_limit=aug_contrast_limit,
p=1.0),
albm.HorizontalFlip()
])
val_transforms = albm.Compose([
albm.Resize(base_size, base_size),
albm.CenterCrop(target_size, target_size)
])
return train_transforms, val_transforms
def __augment_image(self, image):
random.seed()
np.random.seed()
train_aug = A.Compose(
transforms = [
A.RandomSizedCrop(min_max_height=(256, 512), height=512, width=512, p=1),
A.RandomRotate90(p=1),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.IAAAffine(shear=15, p=0.5),
A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=10, p=.3),
A.Blur(blur_limit=2, p=.33),
A.OpticalDistortion(p=.33),
A.RandomBrightnessContrast(p=.33)
],
p=1
)
test_aug = A.Compose(
transforms = [
A.RandomSizedCrop(min_max_height=(256, 512), height=512, width=512, p=1),
A.RandomRotate90(p=1),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
],
p=1
)
if self.mode == 'train':
augmented_image = train_aug(image=image)['image']
else:
augmented_image = test_aug(image=image)['image']
return augmented_image
class ArTDataset(Dataset):
train_tfms = A.Compose([
A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(p=0.2),
A.IAAAffine(scale=(0.9, 1.1),
translate_percent=(-0.1, 0.1),
rotate=(-10, 10),
p=0.3),
])
lambda_tfms = tv.transforms.Compose([
tv.transforms.ToTensor(),
tv.transforms.Lambda(lambda x: 2 * x - 1) # for EffNet advprop
])
def __init__(self,
image_names: str,
image_labels: dict,
is_train=False,
image_size=(768, 768),
scale=4):
self.image_names = image_names
self.image_labels = image_labels
self.is_train = is_train
self.image_size = image_size
self.scale = scale
self._construct_resize_tfms()
def _construct_resize_tfms(self, ):
self.resize_tfms = A.Resize(*self.image_size)
self.resize_mask_tfms = A.Resize(self.image_size[0] // self.scale,
self.image_size[1] // self.scale)
def __len__(self, ):
return len(self.image_names)
def __getitem__(self, idx):
gt_name = self.image_names[idx]
image_org, mask = get_image_and_mask_from_name(gt_name)
# RESIZE
transformed = self.resize_tfms(image=image_org, mask=mask)
image = transformed['image']
mask = transformed['mask']
if self.is_train:
transformed = self.train_tfms(image=image, mask=mask)
image = transformed['image']
mask = transformed['mask']
# NORMALIZE
image = self.lambda_tfms(image)
# RESIZE MASK
transformed = self.resize_mask_tfms(image=image_org, mask=mask)
mask = transformed['mask']
return image, mask
def __init__(self,
images,
labels,
batch_size=16,
image_shape=(256, 512, 1),
do_shuffle_at_epoch_end=True,
length=None,
do_augment=True):
super().__init__(images, labels, batch_size, image_shape,
do_shuffle_at_epoch_end, length, do_augment)
self.augmenting_pipeline = A.Compose([
A.HorizontalFlip(),
A.IAAAffine(translate_percent={"x": (-1, 1)}, mode="reflect", p=1),
A.PadIfNeeded(min_width=int(self.input_shape[1] * 2),
min_height=self.input_shape[0]),
A.GridDistortion(p=0.8, distort_limit=0.5),
A.ElasticTransform(p=0.5,
alpha=10,
sigma=100 * 0.03,
alpha_affine=0),
A.CenterCrop(width=self.input_shape[1],
height=self.input_shape[0]),
A.IAAPerspective(scale=(0, 0.10), p=1),
A.ShiftScaleRotate(shift_limit=0,
scale_limit=(.0, 0.4),
rotate_limit=0,
p=0.5),
A.CLAHE(clip_limit=2.0, p=0.5),
A.Lambda(
image=self.convert_image,
mask=self.convert_segmentations,
),
])
def __init__(self, n, m):
self.n = n
self.m = m
m_ratio = self.m / 30.0
self.augment_list = (
A.CLAHE(always_apply=True),
A.Equalize(always_apply=True),
A.InvertImg(always_apply=True),
A.Rotate(limit=30 * m_ratio, always_apply=True),
A.Posterize(num_bits=int(4 * m_ratio), always_apply=True),
A.Solarize(threshold=m_ratio, always_apply=True),
A.RGBShift(r_shift_limit=110 * m_ratio,
g_shift_limit=110 * m_ratio,
b_shift_limit=110 * m_ratio,
always_apply=True),
A.HueSaturationValue(hue_shift_limit=20 * m_ratio,
sat_shift_limit=30 * m_ratio,
val_shift_limit=20 * m_ratio,
always_apply=True),
A.RandomContrast(limit=m_ratio, always_apply=True),
A.RandomBrightness(limit=m_ratio, always_apply=True),
# A.Sharpen(always_apply=True), 0.1, 1.9),
A.ShiftScaleRotate(shift_limit=0.3 * m_ratio,
shift_limit_y=0,
rotate_limit=0,
always_apply=True),
A.ShiftScaleRotate(shift_limit=0.3 * m_ratio,
shift_limit_x=0,
rotate_limit=0,
always_apply=True),
A.Cutout(num_holes=int(8 * m_ratio), always_apply=True),
A.IAAAffine(shear=0.3 * m_ratio, always_apply=True))
assert self.n <= len(self.augment_list)
def get_aug(aug_type="val", size=512):
"""Return augmentations by type
Args:
aug_type (str): one of `val`, `test`, `light`, `medium`
size (int): final size of the crop
"""
NORM_TO_TENSOR = albu.Compose(
[albu.Normalize(mean=MEAN, std=STD),
albu_pt.ToTensorV2()])
# CROP_AUG = albu.Compose([albu.RandomCrop(size, size)])
VAL_AUG = albu.Compose([
# albu.CenterCrop(size, size),
NORM_TO_TENSOR
])
LIGHT_AUG = albu.Compose(
[
albu.Flip(),
albu.Cutout(num_holes=12,
max_h_size=size // 16,
max_w_size=size // 16,
fill_value=0,
p=0.5), NORM_TO_TENSOR
],
p=1.0,
)
# aug from github.com/lyakaap/Landmark2019-1st-and-3rd-Place-Solution/
HARD_AUG = albu.Compose(
[
albu.Flip(p=0.5),
albu.IAAAffine(rotate=0.2, shear=0.2, mode='constant', p=0.5),
albu.RandomBrightnessContrast(
brightness_limit=0.3, contrast_limit=0.3, p=0.5),
albu.MotionBlur(p=0.5),
albu.CLAHE(p=0.5),
albu.Cutout(num_holes=16,
max_h_size=size // 16,
max_w_size=size // 16,
fill_value=0,
p=0.5),
RandomCropThenScaleToOriginalSize(limit=0.3, p=1.0),
NORM_TO_TENSOR,
],
p=1.0,
)
types = {
"val": VAL_AUG,
"test": VAL_AUG,
"light": LIGHT_AUG,
"hard": HARD_AUG,
}
return types[aug_type]
def augmentation(image_size, train=True):
max_crop = image_size // 5
if train:
data_transform = A.Compose([
A.Resize(image_size, image_size),
A.Compose([
A.OneOf([
A.RandomRain(p=0.1),
A.GaussNoise(mean=15),
A.GaussianBlur(blur_limit=10, p=0.4),
A.MotionBlur(p=0.2)
]),
A.OneOf([
A.RGBShift(p=1.0,
r_shift_limit=(-10, 10),
g_shift_limit=(-10, 10),
b_shift_limit=(-10, 10)),
A.RandomBrightnessContrast(
brightness_limit=0.3, contrast_limit=0.1, p=1),
A.HueSaturationValue(hue_shift_limit=20, p=1),
],
p=0.6),
A.OneOf([
A.CLAHE(clip_limit=2),
A.IAASharpen(),
A.IAAEmboss(),
]),
A.OneOf([A.IAAPerspective(p=0.3),
A.ElasticTransform(p=0.1)]),
A.OneOf([
A.Rotate(limit=25, p=0.6),
A.IAAAffine(
scale=0.9,
translate_px=15,
rotate=25,
shear=0.2,
)
],
p=1),
A.Cutout(num_holes=1,
max_h_size=max_crop,
max_w_size=max_crop,
p=0.2)
],
p=1),
A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
AT.ToTensor()
])
else:
data_transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((image_size, image_size)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
return data_transform
def __call__(self, img):
img = aug(
np.array(img),
albu.IAAAffine(scale=self.scale,
rotate=self.rotate,
shear=self.shear,
mode=self.mode,
p=self.p))
img = Image.fromarray(img)
return img
def __init__(self, resize, mean, std, **args):
self.transform = A.Compose([
A.Resize(resize, resize, p=1.0),
A.HorizontalFlip(p = 0.5),
A.ShiftScaleRotate(p = 0.5),
A.IAAAffine(scale = [0.8, 1.2], rotate = [-10, 10], shear = 10, p = 0.5),
A.OneOf([
A.RandomBrightness(limit = 0.2, p = 0.5),
A.RandomBrightnessContrast(brightness_limit = 0.2, contrast_limit = 0.1, p = 0.5),
], p = 0.5),
A.Normalize(mean = mean, std = std, max_pixel_value = 255.0, p = 1.0),
ToTensorV2(p = 1.0),
])
def aug_transforms(
presize: int = 260,
size: int = 224,
interpolation: int = 1,
hflip: Union[float, bool] = 0.5,
vflip: Union[float, bool] = False,
max_lighting: Union[float, bool] = 0.2,
p_lighting: float = 0.75,
max_rotate: Union[float, int, bool] = 10.0,
p_rotate: float = 0.5,
max_warp: Union[float, bool] = 0.2,
p_affine: float = 0.75,
pad_mode: str = "reflect",
mult: float = 1.0,
xtra_tfms: Optional[List] = None,
) -> A.Compose:
"""
Utility func to easily create a list of flip, rotate, zoom, lighting transforms.
Inspired from : https://docs.fast.ai/vision.augment.html#aug_transforms
"""
max_rotate, max_lighting, max_warp = (
np.array([max_rotate, max_lighting, max_warp]) * mult)
transforms = [
A.Resize(presize,
presize,
interpolation=interpolation,
always_apply=True)
]
if hflip:
transforms += [A.HorizontalFlip(p=hflip)]
if vflip:
transforms += [A.VerticalFlip(p=vflip)]
if max_rotate:
transforms += [
A.Rotate(limit=max_rotate, interpolation=interpolation, p=p_rotate)
]
if max_warp:
transforms += [A.IAAAffine(scale=max_warp, p=p_affine, mode=pad_mode)]
if max_lighting:
transforms += [A.RandomBrightness(max_lighting, p=p_lighting)]
transforms += [A.RandomContrast(max_lighting, p=p_lighting)]
if xtra_tfms is not None:
transforms += [xtra_tfms]
transforms += [
A.RandomResizedCrop(size,
size,
interpolation=interpolation,
always_apply=True)
]
return A.Compose(transforms)
def aug(file):
if file == 'train':
return al.Compose([
al.VerticalFlip(p=0.5),
al.HorizontalFlip(p=0.5),
al.RandomRotate90(p=0.5),
al.OneOf([
al.GaussNoise(0.002, p=0.5),
al.IAAAffine(p=0.5),
],
p=0.2),
al.OneOf([
al.Blur(blur_limit=(3, 10), p=0.4),
al.MedianBlur(blur_limit=3, p=0.3),
al.MotionBlur(p=0.3)
],
p=0.3),
al.OneOf([
al.RandomBrightness(p=0.3),
al.RandomContrast(p=0.4),
al.RandomGamma(p=0.3)
],
p=0.5),
al.Cutout(num_holes=20, max_h_size=20, max_w_size=20, p=0.5),
al.ShiftScaleRotate(shift_limit=0.0625,
scale_limit=(0.9, 1),
rotate_limit=45,
p=0.3),
al.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
p=1),
ToTensorV2(p=1)
])
elif file == 'validation':
return al.Compose([
al.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
p=1),
ToTensorV2(p=1)
])
elif file == 'test':
return al.Compose([
al.Normalize(
mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), p=1),
ToTensorV2(p=1)
],
p=1)
def light_augmentations():
return A.Compose([
A.HorizontalFlip(),
A.RandomBrightnessContrast(),
A.OneOf([
A.ShiftScaleRotate(scale_limit=0.05,
rotate_limit=15,
border_mode=cv2.BORDER_CONSTANT),
A.IAAAffine(),
A.IAAPerspective(),
A.NoOp()
]),
A.HueSaturationValue(),
A.Normalize()
])
def get_yolo_transform(img_size, mode='train'):
if mode == 'train':
scale = 1.1
transform = A.Compose([
A.LongestMaxSize(max_size=int(img_size * scale)),
A.PadIfNeeded(min_height=int(img_size * scale),
min_width=int(img_size * scale),
border_mode=cv2.BORDER_CONSTANT),
A.RandomCrop(width=img_size, height=img_size),
A.ColorJitter(
brightness=0.6, contrast=0.6, saturation=0.6, hue=0.6, p=0.4),
A.OneOf([
A.ShiftScaleRotate(
rotate_limit=10, border_mode=cv2.BORDER_CONSTANT, p=0.4),
A.IAAAffine(shear=10, mode='constant', p=0.4)
],
p=1.0),
A.HorizontalFlip(p=0.5),
A.Blur(p=0.1),
A.CLAHE(p=0.1),
A.Posterize(p=0.1),
A.ToGray(p=0.1),
A.ChannelShuffle(p=0.05),
A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255),
ToTensorV2()
],
bbox_params=A.BboxParams(format='yolo',
min_visibility=0.4,
label_fields=[]))
elif mode == 'test':
transform = A.Compose([
A.LongestMaxSize(max_size=img_size),
A.PadIfNeeded(min_height=img_size,
min_width=img_size,
border_mode=cv2.BORDER_CONSTANT),
A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255),
ToTensorV2(),
],
bbox_params=A.BboxParams(format="yolo",
min_visibility=0.4,
label_fields=[]))
else:
raise ValueError("'mode' can only accept 'train' or 'test'")
return transform
def default_adjusted(image, mask):
alpha = randrange(30, 45)
sigma = randrange(5, 7)
scale = uniform(0.015, 0.075)
train_transform = [
albu.ShiftScaleRotate(scale_limit=0.2,
rotate_limit=0,
shift_limit=0.1,
p=1,
border_mode=0),
albu.PadIfNeeded(min_height=512,
min_width=512,
always_apply=True,
border_mode=0),
albu.IAAAdditiveGaussianNoise(p=0.2),
albu.IAAPerspective(p=0.5),
albu.OneOf(
[
albu.CLAHE(p=1),
albu.RandomBrightnessContrast(p=1),
albu.RandomGamma(p=1),
],
p=0.9,
),
albu.OneOf(
[
albu.RandomBrightnessContrast(brightness_limit=0, p=1),
albu.HueSaturationValue(p=1),
],
p=0.9,
),
albu.OneOf(
[
albu.ElasticTransform(alpha=alpha,
sigma=sigma,
p=1,
border_mode=cv2.BORDER_CONSTANT),
albu.IAAAffine(scale=scale, p=1, mode="constant"),
],
p=0.7,
),
]
aug_func = albu.Compose(train_transform)
augmented = aug_func(image=image, mask=mask)
return augmented['image'], augmented['mask']
def __init__(self, resize, mean, std, **args):
self.transform = A.Compose([
A.Resize(resize, resize, p=1.0),
A.ShiftScaleRotate(p = 0.5),
A.IAAAffine(scale = [0.8, 1.2], rotate = [-10, 10], shear = 10, p = 0.5),
A.GaussNoise(var_limit=(10.0, 50.0), mean=0, always_apply=False, p=0.5),
A.GaussianBlur(blur_limit=(3, 7), sigma_limit=0, always_apply=False, p=0.5),
A.OneOf([
A.MedianBlur(blur_limit=7, always_apply=False, p=1.0),
A.MotionBlur(blur_limit=7, always_apply=False, p=1.0),
], p = 0.5),
A.OneOf([
A.GlassBlur(sigma=0.7, max_delta=4, iterations=2, always_apply=False, mode='fast', p=1.0),
A.Blur(blur_limit=7, always_apply=False, p=1.0),
], p = 0.5),
A.Normalize(mean = mean, std = std, max_pixel_value = 255.0, p = 1.0),
ToTensorV2(p = 1.0),
])
def shape_aug(img, lbl, input_shape=(270, 270)):
assert img.dtype == 'uint8'
assert lbl.dtype == 'uint8'
augumented = albu.Compose([
albu.IAAAffine(scale=(0.8, 1.2),
rotate=179,
shear=5,
translate_percent=(0.01, 0.01),
mode='constant',
p=0.5),
albu.Flip(p=.5),
albu.CenterCrop(input_shape[0], input_shape[1], always_apply=True)
])
augs = augumented(image=img, mask=lbl)
img, lbl = augs['image'], augs['mask']
return img, lbl
def parse_annotation(annotation, train_input_size, annotation_type):
line = annotation.split()
image_path = line[0]
image = np.array(cv2.imread(image_path))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
exist_boxes = True
augmentation = al.Compose([
# al.VerticalFlip(p=0.5),
# al.HorizontalFlip(p=0.5),
# al.RandomRotate90(p=0.5),
al.OneOf([
al.GaussNoise(0.002, p=0.5),
al.IAAAffine(p=0.5),
], p=0.2),
al.OneOf([
al.Blur(blur_limit=(3, 10), p=0.4),
al.MedianBlur(blur_limit=3, p=0.3),
al.MotionBlur(p=0.3)
],
p=0.3),
al.OneOf([
al.RandomBrightness(p=0.3),
al.RandomContrast(p=0.4),
al.RandomGamma(p=0.3)
],
p=0.5),
al.Cutout(num_holes=20, max_h_size=20, max_w_size=20, p=0.5),
])
bboxes = np.array([
list(map(lambda x: int(float(x)), box.split(','))) for box in line[1:]
])
if annotation_type == 'train':
# image, bboxes = random_horizontal_flip(np.copy(image), np.copy(bboxes))
# image, bboxes = random_crop(np.copy(image), np.copy(bboxes))
# image, bboxes = random_translate(np.copy(image), np.copy(bboxes))
augm = augmentation(image=image)
image = augm['image'].astype(np.float32) / 255.
# image, bboxes = image_preprocess(np.copy(image), [train_input_size, train_input_size], np.copy(bboxes))
return image, bboxes, exist_boxes
def clf_train_augs(height: int, width: int) -> albu.Compose:
"""Training augmentations for classification. We prefer for this model to be really
robust. Feel free to tweak these paramters or ad other augmentations."""
return albu.Compose([
albu.Resize(height=height, width=width),
albu.OneOf([
albu.IAAAffine(shear=6, rotate=5, always_apply=True),
albu.ShiftScaleRotate(shift_limit=0.025,
scale_limit=0.1,
rotate_limit=10),
]),
albu.ShiftScaleRotate(shift_limit=0.025,
scale_limit=0.1,
rotate_limit=10),
albu.Flip(),
albu.RandomRotate90(),
albu.OneOf(
[
albu.HueSaturationValue(p=1.0),
albu.IAAAdditiveGaussianNoise(p=1.0),
albu.IAASharpen(p=1.0),
albu.RandomBrightnessContrast(
brightness_limit=0.1, contrast_limit=0.1, p=1.0),
albu.RandomGamma(p=1.0),
],
p=1.0,
),
albu.OneOf(
[
albu.Blur(blur_limit=3, p=1.0),
albu.MedianBlur(blur_limit=3, p=1.0),
albu.MotionBlur(blur_limit=3, p=1.0),
],
p=1.0,
),
albu.Normalize(),
])
def composeAugmentation(self):
if (source == "train" :)
self.augment = albumentations.Compose(
[
albumentations.HorizontalFlip(),
albumentations.Rotate(10, always_apply=True),
albumentations.SmallestMaxSize(self.sizeX+8, always_apply=True),
albumentations.CenterCrop(self.sizeX, self.sizeY, always_apply=True),
albumentations.GridDistortion(always_apply=False),
albumentations.IAAAffine(rotate=2, shear=5, always_apply=False),
#albumentations.OpticalDistortion(),
albumentations.ElasticTransform(alpha=128, sigma=64, always_apply=True, alpha_affine=0),
albumentations.RandomBrightnessContrast(0.1, 0.1, always_apply=True),
]
)
if (source == "val"):
self.augment = albumentations.Compose(
[
albumentations.SmallestMaxSize(self.sizeX+8, always_apply=True),
albumentations.CenterCrop(self.sizeX, self.sizeY, always_apply=True),
]
)
def train_function(gpu, world_size, node_rank, gpus):
import torch.multiprocessing
torch.multiprocessing.set_sharing_strategy('file_system')
torch.manual_seed(25)
np.random.seed(25)
rank = node_rank * gpus + gpu
dist.init_process_group(
backend='nccl',
init_method='env://',
world_size=world_size,
rank=rank
)
width_size = 512
batch_size = 32
accumulation_step = 5
device = torch.device("cuda:{}".format(gpu) if torch.cuda.is_available() else "cpu")
if rank == 0:
wandb.init(project='inception_v3', group=wandb.util.generate_id())
wandb.config.width_size = width_size
wandb.config.aspect_rate = 1
wandb.config.batch_size = batch_size
wandb.config.accumulation_step = accumulation_step
shutil.rmtree('tensorboard_runs', ignore_errors=True)
writer = SummaryWriter(log_dir='tensorboard_runs', filename_suffix=str(time.time()))
ranzcr_df = pd.read_csv('train_folds.csv')
ranzcr_train_df = ranzcr_df[ranzcr_df['fold'] != 1]
chestx_df = pd.read_csv('chestx_pseudolabeled_data_lazy_balancing.csv')
train_image_transforms = alb.Compose([
alb.ImageCompression(quality_lower=65, p=0.5),
alb.HorizontalFlip(p=0.5),
alb.CLAHE(p=0.5),
alb.OneOf([
alb.GridDistortion(
num_steps=8,
distort_limit=0.5,
p=1.0
),
alb.OpticalDistortion(
distort_limit=0.5,
shift_limit=0.5,
p=1.0,
),
alb.ElasticTransform(alpha=3, p=1.0)],
p=0.7
),
alb.RandomResizedCrop(
height=width_size,
width=width_size,
scale=(0.8, 1.2),
p=0.7
),
alb.RGBShift(p=0.5),
alb.RandomSunFlare(p=0.5),
alb.RandomFog(p=0.5),
alb.RandomBrightnessContrast(p=0.5),
alb.HueSaturationValue(
hue_shift_limit=20,
sat_shift_limit=20,
val_shift_limit=20,
p=0.5
),
alb.ShiftScaleRotate(shift_limit=0.025, scale_limit=0.1, rotate_limit=20, p=0.5),
alb.CoarseDropout(
max_holes=12,
min_holes=6,
max_height=int(width_size / 6),
max_width=int(width_size / 6),
min_height=int(width_size / 6),
min_width=int(width_size / 20),
p=0.5
),
alb.IAAAdditiveGaussianNoise(loc=0, scale=(2.5500000000000003, 12.75), per_channel=False, p=0.5),
alb.IAAAffine(scale=1.0, translate_percent=None, translate_px=None, rotate=0.0, shear=0.0, order=1, cval=0,
mode='reflect', p=0.5),
alb.IAAAffine(rotate=90., p=0.5),
alb.IAAAffine(rotate=180., p=0.5),
alb.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2()
])
train_set = NoisyStudentDataset(ranzcr_train_df, chestx_df, train_image_transforms,
'../ranzcr/train', '../data', width_size=width_size)
train_sampler = DistributedSampler(train_set, num_replicas=world_size, rank=rank, shuffle=True)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=False, num_workers=4, sampler=train_sampler)
ranzcr_valid_df = ranzcr_df[ranzcr_df['fold'] == 1]
valid_image_transforms = alb.Compose([
alb.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2()
])
valid_set = ImageDataset(ranzcr_valid_df, valid_image_transforms, '../ranzcr/train', width_size=width_size)
valid_loader = DataLoader(valid_set, batch_size=batch_size, num_workers=4, pin_memory=False, drop_last=False)
# ranzcr_valid_df = ranzcr_df[ranzcr_df['fold'] == 1]
# valid_image_transforms = alb.Compose([
# alb.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
# ToTensorV2()
# ])
# valid_set = ImageDataset(ranzcr_valid_df, valid_image_transforms, '../ranzcr/train', width_size=width_size)
# valid_sampler = DistributedSampler(valid_set, num_replicas=world_size, rank=rank)
# valid_loader = DataLoader(valid_set, batch_size=batch_size, num_workers=4, sampler=valid_sampler)
checkpoints_dir_name = 'inception_v3_noisy_student_{}'.format(width_size)
os.makedirs(checkpoints_dir_name, exist_ok=True)
# model = EfficientNetNoisyStudent(11, pretrained_backbone=True,
# mixed_precision=True, model_name='tf_efficientnet_b7_ns')
model = Inception(11, pretrained_backbone=True, mixed_precision=False, model_name='inception_v3')
model = SyncBatchNorm.convert_sync_batchnorm(model)
model.to(device)
model = DistributedDataParallel(model, device_ids=[gpu])
# class_weights = [354.625, 23.73913043478261, 2.777105767812362, 110.32608695652173,
# 52.679245283018865, 9.152656621728786, 4.7851333032083145,
# 8.437891632878731, 2.4620064899945917, 0.4034751151063363, 31.534942820838626]
class_names = ['ETT - Abnormal', 'ETT - Borderline', 'ETT - Normal',
'NGT - Abnormal', 'NGT - Borderline', 'NGT - Incompletely Imaged', 'NGT - Normal',
'CVC - Abnormal', 'CVC - Borderline', 'CVC - Normal', 'Swan Ganz Catheter Present']
scaler = GradScaler()
criterion = torch.nn.BCEWithLogitsLoss()
lr_start = 1e-4
lr_end = 1e-6
weight_decay = 0
epoch_num = 20
if rank == 0:
wandb.config.model_name = checkpoints_dir_name
wandb.config.lr_start = lr_start
wandb.config.lr_end = lr_end
wandb.config.weight_decay = weight_decay
wandb.config.epoch_num = epoch_num
wandb.config.optimizer = 'adam'
wandb.config.scheduler = 'CosineAnnealingLR'
wandb.config.is_loss_weights = 'no'
optimizer = Adam(model.parameters(), lr=lr_start, weight_decay=weight_decay)
scheduler = CosineAnnealingLR(optimizer, T_max=epoch_num, eta_min=lr_end, last_epoch=-1)
max_val_auc = 0
for epoch in range(epoch_num):
train_loss, train_avg_auc, train_auc, train_rocs, train_data_pr, train_duration = one_epoch_train(
model, train_loader, optimizer, criterion, device, scaler,
iters_to_accumulate=accumulation_step, clip_grads=False)
scheduler.step()
if rank == 0:
val_loss, val_avg_auc, val_auc, val_rocs, val_data_pr, val_duration = eval_model(
model, valid_loader, device, criterion, scaler)
wandb.log({'train_loss': train_loss, 'val_loss': val_loss,
'train_auc': train_avg_auc, 'val_auc': val_avg_auc, 'epoch': epoch})
for class_name, auc1, auc2 in zip(class_names, train_auc, val_auc):
wandb.log({'{} train auc'.format(class_name): auc1,
'{} val auc'.format(class_name): auc2, 'epoch': epoch})
if val_avg_auc > max_val_auc:
max_val_auc = val_avg_auc
wandb.run.summary["best_accuracy"] = val_avg_auc
print('EPOCH %d:\tTRAIN [duration %.3f sec, loss: %.3f, avg auc: %.3f]\t\t'
'VAL [duration %.3f sec, loss: %.3f, avg auc: %.3f]\tCurrent time %s' %
(epoch + 1, train_duration, train_loss, train_avg_auc,
val_duration, val_loss, val_avg_auc, str(datetime.now(timezone('Europe/Moscow')))))
torch.save(model.module.state_dict(),
os.path.join(checkpoints_dir_name, '{}_epoch{}_val_auc{}_loss{}_train_auc{}_loss{}.pth'.format(
checkpoints_dir_name, epoch + 1, round(val_avg_auc, 3), round(val_loss, 3),
round(train_avg_auc, 3), round(train_loss, 3))))
if rank == 0:
wandb.finish()
def val_albu_augment(record):
"""
"""
verbose = record.get('verbose', False)
image = record['image']
mask = record['mask']
if verbose:
pipeline = albu.ReplayCompose
else:
pipeline = albu.Compose
aug = pipeline([
albu.OneOf([
albu.RandomBrightnessContrast(brightness_limit=0.2,
contrast_limit=0.2,
brightness_by_max=True,
always_apply=False,
p=1),
albu.RandomBrightnessContrast(brightness_limit=(-0.2, 0.6),
contrast_limit=.2,
brightness_by_max=True,
always_apply=False,
p=1),
albu.augmentations.transforms.ColorJitter(brightness=0.2,
contrast=0.2,
saturation=0.1,
hue=0.1,
always_apply=False,
p=1),
albu.RandomGamma(p=1)
],
p=1),
albu.OneOf([
albu.GaussNoise(0.02, p=.5),
albu.IAAAffine(p=.5),
],
p=.25),
albu.OneOf([
albu.augmentations.transforms.Blur(
blur_limit=15, always_apply=False, p=0.25),
albu.augmentations.transforms.Blur(
blur_limit=3, always_apply=False, p=0.5)
],
p=0.5),
])
data = aug(image=image, mask=mask)
record['image'] = data['image']
record['mask'] = data['mask']
if verbose:
for transformation in data['replay']['transforms']:
if not isinstance(transformation, dict):
print('not a dict')
pass
elif transformation.get('applied', False):
print(30 * '-')
if 'OneOf' in transformation['__class_fullname__']:
print(30 * '=')
for _trans in transformation['transforms']:
if not _trans.get('applied', False): continue
_name = _trans['__class_fullname__']
if 'Flip' in _name: continue
print(_trans['__class_fullname__'])
for k, v in _trans.items():
if k in [
'__class_fullname__', 'applied',
'always_apply'
]:
continue
print(f"{k}: {v}")
else:
_name = transformation['__class_fullname__']
if 'Flip' in _name: continue
print(_name)
for k, v in transformation.items():
if k in [
'__class_fullname__', 'applied', 'always_apply'
]:
continue
print(f"{k}: {v}")
return record
car_name = "classes_carros"
plt_folder = "*/classes"
car_folder = "*/classes_carros"
ocr_file = '%s/OCR/output.txt' % (path)
metadata_length = 35
tam_max = 3
L1_layer = Lambda(lambda tensor: K.abs(tensor[0] - tensor[1]))
train_augs = [[], [], [], [], [], []]
test_augs = [[], [], [], []]
keys = ['Set01', 'Set02', 'Set03', 'Set04', 'Set05']
seq_car = albu.Compose([
albu.IAACropAndPad(px=(0, 8)),
albu.IAAAffine(scale=(0.8, 1.2),
shear=(-8, 8),
order=[0, 1],
cval=(0),
mode='constant'),
albu.ToFloat(max_value=255)
],
p=0.7)
seq_car2 = albu.Compose([
albu.IAACropAndPad(px=(0, 8)),
albu.IAAAffine(scale=(0.8, 1.2),
shear=(-8, 8),
order=[0, 1],
cval=(0),
mode='constant'),
],
p=0.7)
train_transforms = A.Compose(
[
A.LongestMaxSize(max_size=int(IMAGE_SIZE * scale)),
A.PadIfNeeded(
min_height=int(IMAGE_SIZE * scale),
min_width=int(IMAGE_SIZE * scale),
border_mode=cv2.BORDER_CONSTANT,
),
A.RandomCrop(width=IMAGE_SIZE, height=IMAGE_SIZE),
A.ColorJitter(
brightness=0.6, contrast=0.6, saturation=0.6, hue=0.6, p=0.4),
A.OneOf(
[
A.ShiftScaleRotate(
rotate_limit=10, p=0.4, border_mode=cv2.BORDER_CONSTANT),
A.IAAAffine(shear=10, p=0.4, mode="constant"),
],
p=1.0,
),
A.HorizontalFlip(p=0.5),
A.Blur(p=0.1),
A.CLAHE(p=0.1),
A.Posterize(p=0.1),
A.ToGray(p=0.1),
A.ChannelShuffle(p=0.05),
A.Normalize(
mean=[0, 0, 0],
std=[1, 1, 1],
max_pixel_value=255,
),
ToTensorV2(),
def run(fold, args):
if args.sz:
print(f"Images will be resized to {args.sz}")
args.sz = int(args.sz)
# get training and valid data
df = pd.read_csv(args.training_folds_csv)
if args.loss == 'crossentropy' and not args.isic2019:
diag_to_ix = {
v: i
for i, v in enumerate(sorted(list(set(df.diagnosis))))
}
ix_to_diag = {v: i for i, v in diag_to_ix.items()}
if args.external_csv_path:
df_external = pd.read_csv(args.external_csv_path)
df_train = df.query(f"kfold != {fold}").reset_index(drop=True)
df_valid = df.query(f"kfold == {fold}").reset_index(drop=True)
print(
f"Running for K-Fold {fold}; train_df: {df_train.shape}, valid_df: {df_valid.shape}"
)
# calculate weights for NN loss
weights = len(df) / df.target.value_counts().values
class_weights = torch.FloatTensor(weights)
if args.loss == 'weighted_bce':
print(f"assigning weights {weights} to loss fn.")
if args.loss == 'focal_loss':
print("Focal loss will be used for training.")
if args.loss == 'weighted_cross_entropy':
print(f"assigning weights {weights} to loss fn.")
# create model
if 'efficient_net' in args.model_name:
model = MODEL_DISPATCHER[args.model_name](
pretrained=args.pretrained,
arch_name=args.arch_name,
ce=(args.loss == 'crossentropy'
or args.loss == 'weighted_cross_entropy'
or args.load_pretrained_2019))
else:
model = MODEL_DISPATCHER[args.model_name](pretrained=args.pretrained)
if args.model_path is not None:
print(
f"Loading pretrained model and updating final layer from {args.model_path}"
)
model.load_state_dict(torch.load(args.model_path))
nftrs = model.base_model._fc.in_features
model.base_model._fc = nn.Linear(nftrs, 1)
meta_array = None
if args.use_metadata:
# create meta array
sex_dummy_train = pd.get_dummies(df_train['sex'])[['male', 'female']]
site_dummy_train = pd.get_dummies(
df_train['anatom_site_general_challenge'])[[
'head/neck', 'lower extremity', 'oral/genital', 'palms/soles',
'torso', 'upper extremity'
]]
assert max(df_train.age_approx) < 100
age_train = df_train.age_approx.fillna(-5) / 100
meta_array = pd.concat([sex_dummy_train, site_dummy_train, age_train],
axis=1).values
# modify model forward
if args.freeze_cnn:
model.load_state_dict(torch.load(args.model_path))
# update the forward pass
model = modify_model(model, args)
# freeze cnn
if args.freeze_cnn:
print("\nFreezing CNN layers!\n")
for param in model.base_model.parameters():
param.requires_grad = False
# add external meta to meta array
if args.external_csv_path:
sex_dummy_ext = pd.get_dummies(
df_external['sex'])[['male', 'female']]
df_external[
'anatom_site_general'] = df_external.anatom_site_general.replace(
{
'anterior torso': 'torso',
'lateral torso': 'torso',
'posterior torso': 'torso'
})
site_dummy_ext = pd.get_dummies(
df_external['anatom_site_general'])[[
'head/neck', 'lower extremity', 'oral/genital',
'palms/soles', 'torso', 'upper extremity'
]]
assert max(df_external.age_approx) < 100
age_ext = df_external.age_approx.fillna(-5) / 100
meta_array = np.concatenate([
meta_array,
pd.concat([sex_dummy_ext, site_dummy_ext, age_ext],
axis=1).values
])
assert meta_array.shape[1] == 9
model = model.to(args.device)
train_aug = albumentations.Compose([
albumentations.RandomScale(0.07),
albumentations.Rotate(50),
albumentations.RandomBrightnessContrast(0.15, 0.1),
albumentations.Flip(p=0.5),
albumentations.IAAAffine(shear=0.1),
albumentations.RandomCrop(args.sz, args.sz)
if args.sz else albumentations.NoOp(),
albumentations.OneOf([
albumentations.Cutout(random.randint(1, 8), 16, 16),
albumentations.CoarseDropout(random.randint(1, 8), 16, 16)
]),
albumentations.Normalize(always_apply=True)
])
valid_aug = albumentations.Compose([
albumentations.CenterCrop(args.sz, args.sz)
if args.sz else albumentations.NoOp(),
albumentations.Normalize(always_apply=True),
])
print(f"\nUsing train augmentations: {train_aug}\n")
# get train and valid images & targets and add external data if required (external data only contains melonama data)
train_images = df_train.image_name.tolist()
if args.external_csv_path:
external_images = df_external.image.tolist()
if args.exclude_outliers_2019:
# from EDA notebook
external_images = np.load(
f'/home/ubuntu/repos/kaggle/melonama/data/external/clean_external_2019_{args.sz}.npy'
).tolist()
print(
f"\n\n{len(external_images)} external images will be added to each training fold."
)
train_images = train_images + external_images
if args.use_pseudo_labels:
test_df = pd.read_csv(
'/home/ubuntu/repos/kaggle/melonama/data/test.csv')
test_images = test_df.image_name.tolist()
if args.pseudo_images_path:
test_images = list(
np.load(args.pseudo_images_path, allow_pickle=True))
print(
f"\n\n{len(test_images)} test images will be added to each training fold."
)
train_images = train_images + test_images
train_image_paths = [
os.path.join(args.train_data_dir, image_name + '.jpg')
for image_name in train_images
]
train_targets = df_train.target if not args.external_csv_path else np.concatenate(
[df_train.target.values,
np.ones(len(external_images))])
if args.use_pseudo_labels:
train_targets = np.concatenate([
train_targets,
np.load(args.pseudo_labels_path, allow_pickle=True)
])
if args.loss == 'crossentropy':
df_train['diagnosis'] = df_train.diagnosis.map(diag_to_ix)
train_targets = df_train.diagnosis.values
assert len(train_image_paths) == len(
train_targets
), "Length of train images {} doesnt match length of targets {}".format(
len(train_images), len(train_targets))
# same for valid dataframe
valid_images = df_valid.image_name.tolist()
valid_image_paths = [
os.path.join(args.train_data_dir, image_name + '.jpg')
for image_name in valid_images
]
valid_targets = df_valid.target
if args.loss == 'crossentropy':
df_valid['diagnosis'] = df_valid.diagnosis.map(diag_to_ix)
valid_targets = df_valid.diagnosis.values
print(
f"\n\n Total Train images: {len(train_image_paths)}, Total val: {len(valid_image_paths)}\n\n"
)
# create train and valid dataset, dont use color constancy as already preprocessed in directory
train_dataset = MelonamaDataset(train_image_paths,
train_targets,
train_aug,
cc=args.cc,
meta_array=meta_array)
valid_dataset = MelonamaDataset(valid_image_paths,
valid_targets,
valid_aug,
cc=args.cc,
meta_array=meta_array)
# create dataloaders
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.train_batch_size,
shuffle=True,
num_workers=4)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=args.valid_batch_size,
shuffle=False,
num_workers=4)
# create optimizer and scheduler for training
optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=[3, 5, 6, 7, 8, 9, 10, 11, 13, 15], gamma=0.5)
es = EarlyStopping(patience=3,
mode='min' if args.metric == 'valid_loss' else 'max')
for epoch in range(args.epochs):
train_loss = train_one_epoch(
args,
train_loader,
model,
optimizer,
weights=None
if not args.loss.startswith('weighted') else class_weights)
preds, valid_loss = evaluate(args, valid_loader, model)
predictions = np.vstack(preds).ravel()
if args.loss == 'crossentropy' or args.loss == 'weighted_cross_entropy':
accuracy = metrics.accuracy_score(valid_targets, predictions)
else:
auc = metrics.roc_auc_score(valid_targets, predictions)
preds_df = pd.DataFrame({
'predictions': predictions,
'targets': valid_targets,
'valid_image_paths': valid_image_paths
})
print(
f"Epoch: {epoch}, Train loss: {train_loss}, Valid loss: {valid_loss}, Valid Score: {locals()[f'{args.metric}']}"
)
scheduler.step()
for param_group in optimizer.param_groups:
print(f"Current Learning Rate: {param_group['lr']}")
es(locals()[f"{args.metric}"],
model,
model_path=
f"/home/ubuntu/repos/kaggle/melonama/models/{syd_now.strftime(r'%d%m%y')}/{args.arch_name}_fold_{fold}_{args.sz}_{locals()[f'{args.metric}']}.bin",
preds_df=preds_df,
df_path=
f"/home/ubuntu/repos/kaggle/melonama/valid_preds/{syd_now.strftime(r'%d%m%y')}/{args.arch_name}_fold_{fold}_{args.sz}_{locals()[f'{args.metric}']}.bin",
args=args)
if es.early_stop:
return preds_df
train_transforms = A.Compose(
[
A.LongestMaxSize(max_size=int(IMAGE_SIZE * scale)),
A.PadIfNeeded(
min_height=int(IMAGE_SIZE * scale),
min_width=int(IMAGE_SIZE * scale),
border_mode=cv2.BORDER_CONSTANT,
),
A.RandomCrop(width=IMAGE_SIZE, height=IMAGE_SIZE),
A.ColorJitter(
brightness=0.6, contrast=0.6, saturation=0.6, hue=0.6, p=0.4),
A.OneOf(
[
A.ShiftScaleRotate(
rotate_limit=20, p=0.5, border_mode=cv2.BORDER_CONSTANT),
A.IAAAffine(shear=15, p=0.5, mode="constant"),
],
p=1.0,
),
A.HorizontalFlip(p=0.5),
A.Blur(p=0.1),
A.CLAHE(p=0.1),
A.Posterize(p=0.1),
A.ToGray(p=0.1),
A.ChannelShuffle(p=0.05),
A.Normalize(
mean=[0, 0, 0],
std=[1, 1, 1],
max_pixel_value=255,
),
ToTensorV2(),
PIN_MEMORY = True
SAVE_MODEL = True
LOAD_MODEL = True
# Data augmentation for images
train_transforms = A.Compose([
A.Resize(width=760, height=760),
A.RandomCrop(height=728, width=728),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.RandomRotate90(p=0.5),
A.Blur(p=0.3),
A.CLAHE(p=0.3),
A.ColorJitter(p=0.3),
A.CoarseDropout(max_holes=12, max_height=20, max_width=20, p=0.3),
A.IAAAffine(shear=30, rotate=0, p=0.2, mode="constant"),
A.Normalize(
mean=[0.3199, 0.2240, 0.1609],
std=[0.3020, 0.2183, 0.1741],
max_pixel_value=255.0,
),
ToTensorV2(),
])
val_transforms = A.Compose([
A.Resize(height=728, width=728),
A.Normalize(
mean=[0.3199, 0.2240, 0.1609],
std=[0.3020, 0.2183, 0.1741],
max_pixel_value=255.0,
),
MODEL_PATH = 'Christof/models/ResNet18/1_ext/'
exp_suffix = ''
SIZE = 512
# Load dataset info
path_to_train = 'Christof/assets/train_rgb_512/'
data = pd.read_csv('Christof/assets/train.csv')
normal_aug = A.Compose([#A.Rotate((0,30),p=0.75),
A.RandomRotate90(p=1),
A.HorizontalFlip(p=0.5),
#A.RandomBrightness(0.05),
#A.RandomContrast(0.05),
A.IAAAffine(translate_percent=10,rotate=45,shear=10, scale=(0.9,1.1)),
#A.RandomAffine(degrees=45, translate=(0.1,0.1), shear=10, scale=(0.9,1.1))
A.Normalize(mean=(0.08069, 0.05258, 0.05487), std=(0.1300, 0.0879, 0.1386),
max_pixel_value=255.)
])
normal_aug_ext = A.Compose([#A.Rotate((0,30),p=0.75),
A.RandomRotate90(p=1),
A.HorizontalFlip(p=0.5),
#A.RandomBrightness(0.05),
#A.RandomContrast(0.05),
A.IAAAffine(translate_percent=10,rotate=45,shear=10, scale=(0.9,1.1)),
#A.RandomAffine(degrees=45, translate=(0.1,0.1), shear=10, scale=(0.9,1.1))
A.Normalize(mean=(0.1174382, 0.06798691, 0.06592218), std=(0.16392466 ,0.10036821, 0.16703453),
max_pixel_value=255.)
############################## Prapare Augmentation
train_transform = albumentations.Compose([
# albumentations.OneOf([
# albumentations.Resize(IMAGE_HEIGHT_RESIZE, IMAGE_WIDTH_RESIZE),
# albumentations.RandomResizedCrop(IMAGE_HEIGHT_RESIZE, IMAGE_WIDTH_RESIZE, scale=(0.75, 1.0)),
# ], p=1),
albumentations.Resize(IMAGE_HEIGHT_RESIZE, IMAGE_WIDTH_RESIZE),
albumentations.OneOf(
[
# albumentations.Cutout(num_holes=8, max_h_size=2, max_w_size=4, fill_value=0),
GridMask(num_grid=(3, 7), p=1),
albumentations.ShiftScaleRotate(scale_limit=.15,
rotate_limit=20,
border_mode=cv2.BORDER_CONSTANT),
albumentations.IAAAffine(shear=20, mode='constant'),
albumentations.IAAPerspective(),
# albumentations.GridDistortion(distort_limit=0.01),
],
p=0.8)
])
test_transform = albumentations.Compose([
albumentations.Resize(IMAGE_HEIGHT_RESIZE, IMAGE_WIDTH_RESIZE),
])
############################################ Define Dataset
class bengaliai_Dataset(torch.utils.data.Dataset):
def __init__(self, \
data_path, \
def trainval(exp_dict, savedir_base, datadir, reset=False, num_workers=0):
# bookkeepting stuff
# ==================
savedir = os.path.join(savedir_base, hu.hash_dict(exp_dict))
os.makedirs(savedir, exist_ok=True)
if reset:
hc.delete_and_backup_experiment(savedir)
print("Experiment saved in %s" % savedir)
# Dataset
# ==================
# train set
data_transform = A.Compose(
[
A.Flip(p=0.3),
A.IAAAffine(p=0.3),
A.Rotate(p=0.3),
A.HueSaturationValue(hue_shift_limit=10,
sat_shift_limit=15,
val_shift_limit=10,
p=0.3),
A.GaussianBlur(3, p=0.3),
A.GaussNoise(30, p=0.3)
],
keypoint_params=A.KeypointParams(format='xy'),
additional_targets={
'mask0': 'mask',
'mask1': 'mask',
'mask2': 'mask',
'keypoints0': 'keypoints',
'keypoints1': 'keypoints',
'keypoints2': 'keypoints',
'keypoints3': 'keypoints',
'keypoints4': 'keypoints',
'keypoints5': 'keypoints'
})
# random.seed(20201009)
random_seed = random.randint(0, 20201009)
train_set = HEDataset_Fast(data_dir=datadir,
n_classes=exp_dict["n_classes"],
transform=data_transform,
option="Train",
random_seed=random_seed,
obj_option=exp_dict["obj"],
patch_size=exp_dict["patch_size"],
bkg_option=exp_dict["bkg"])
test_transform = A.Compose([A.Resize(1024, 1024)],
keypoint_params=A.KeypointParams(format='xy'),
additional_targets={
'mask0': 'mask',
'mask1': 'mask'
})
# val set
val_set = HEDataset(data_dir=datadir,
transform=test_transform,
option="Validation")
val_loader = DataLoader(val_set, batch_size=1, num_workers=num_workers)
# test set
test_set = HEDataset(data_dir=datadir,
transform=test_transform,
option="Test")
test_loader = DataLoader(test_set, batch_size=1, num_workers=num_workers)
# Model
# ==================
# torch.manual_seed(20201009)
model = models.get_model(exp_dict['model'],
exp_dict=exp_dict,
train_set=train_set).cuda()
model_path = os.path.join(savedir, "model.pth")
score_list_path = os.path.join(savedir, "score_list.pkl")
if os.path.exists(score_list_path):
# resume experiment
model.load_state_dict(hu.torch_load(model_path))
score_list = hu.load_pkl(score_list_path)
s_epoch = score_list[-1]['epoch'] + 1
else:
# restart experiment
score_list = []
s_epoch = 0
# Train & Val
# ==================
print("Starting experiment at epoch %d" % (s_epoch))
# train_sampler = torch.utils.data.RandomSampler(
# train_set, replacement=True, num_samples=2*len(val_set))
train_loader = DataLoader(train_set,
batch_size=exp_dict["batch_size"],
shuffle=True,
num_workers=num_workers)
for e in range(s_epoch, exp_dict['max_epoch']):
# Validate only at the start of each cycle
score_dict = {}
# Train the model
train_dict = model.train_on_loader(train_loader)
# Validate and Visualize the model
val_dict = model.val_on_loader(val_loader,
savedir_images=os.path.join(
savedir, "images"),
n_images=7)
score_dict.update(val_dict)
# Get new score_dict
score_dict.update(train_dict)
score_dict["epoch"] = len(score_list)
# Add to score_list and save checkpoint
score_list += [score_dict]
# Report & Save
score_df = pd.DataFrame(score_list)
print("\n", score_df.tail(), "\n")
hu.torch_save(model_path, model.get_state_dict())
hu.save_pkl(score_list_path, score_list)
print("Checkpoint Saved: %s" % savedir)
# Save Best Checkpoint
if e == 0 or (score_dict.get("val_score", 0) >
score_df["val_score"][:-1].fillna(0).max()):
hu.save_pkl(os.path.join(savedir, "score_list_best.pkl"),
score_list)
hu.torch_save(os.path.join(savedir, "model_best.pth"),
model.get_state_dict())
print("Saved Best: %s" % savedir)
# if s_epoch==exp_dict['max_epoch']:
# e = s_epoch
model.load_state_dict(
hu.torch_load(os.path.join(savedir, "model_best.pth")))
test_dict = model.test_on_loader(test_loader)
hu.save_pkl(os.path.join(savedir, 'test_iou.pkl'), test_dict)
print('Test IoU:{}'.format(test_dict["test_iou"]))
print('Experiment completed et epoch %d' % e)
