Solution for submission 152443

gpu_info = !nvidia-smi
gpu_info = '\n'.join(gpu_info)
if gpu_info.find('failed') >= 0:
  print('Select the Runtime > "Change runtime type" menu to enable a GPU accelerator, ')
  print('and then re-execute this cell.')
else:
  print(gpu_info)

# from google.colab import drive
# drive.mount('/content/drive')

# !pip install -U git+https://github.com/albu/albumentations --no-cache-dir
# !pip install -U segmentation-models-pytorch albumentations --user
# !pip install aicrowd-cli
# !pip install git+https://github.com/qubvel/segmentation_models.pytorch

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import shutil
import cv2
import random
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from tqdm import tqdm
import albumentations as albu
from PIL import Image, ImageEnhance

from torch.utils.data import Dataset, DataLoader
from natsort import natsorted

import torch
import segmentation_models_pytorch as smp

DATA_DIR = '/content/drive/MyDrive/tree_segmentation/'

x_train_dir = os.path.join(DATA_DIR, 'train')
y_train_dir = os.path.join(DATA_DIR, 'trainannot')

x_valid_dir = os.path.join(DATA_DIR, 'val')
y_valid_dir = os.path.join(DATA_DIR, 'valannot')

x_test_dir = os.path.join(DATA_DIR, 'test')
x_test_dir

# from PIL import Image, ImageEnhance
# image,mask = data4viz[return_index(4389)]
# _,g,_ = cv2.split(image)
# print(np.mean(g))
# enhancer = ImageEnhance.Brightness(Image.fromarray(image))
# for i in range(1,6,1):
#   factor = i
#   im_output = enhancer.enhance(factor)
#   im_output = np.array(im_output)
#   _,g,_ = cv2.split(im_output)
#   print(np.mean(g))
#   if np.mean(g)>100:
#     break
# plt.imshow(np.hstack((im_output,image)))
# plt.show()
# plt.imshow(mask.reshape(512,512),cmap='gray')

class TreeSegmentationDataset(Dataset):
  def __init__(self,
               image_directory        = None,
               segmentation_directory = None,
               train                  = True,
               augmentation           = None,
               preprocessing          = None,
               ):
    '''
        We input the directory where segmentations and their corresponding images are present.
        We also consider the augmentation object to be considered while training as well as preprocessing to be done.
    '''
    self.ids                    = os.listdir(image_directory)
    self.image_directory        = image_directory
    self.segmentation_directory = segmentation_directory
    self.augmentation           = augmentation
    self.preprocessing          = preprocessing

    if self.image_directory!=None:
      
      self.image_list           = natsorted(os.listdir(self.image_directory))
      self.segmentation_list    = natsorted(os.listdir(self.segmentation_directory))

    self.train                  = train

  def __getitem__(self,index):
    '''
        Fetches the image corresponding to a particular index from the data.
    '''

    image_location             = self.image_directory        + '/' + self.image_list[index]
    

    image_needed               = cv2.imread(image_location)
    image                      = cv2.cvtColor(image_needed, cv2.COLOR_BGR2RGB)
    
    _,g,_                      = cv2.split(image)
    enhancer                   = ImageEnhance.Brightness(Image.fromarray(image))
    for i in range(1,6,1):
      factor                   = i
      im_output                = enhancer.enhance(factor)
      im_output                = np.array(im_output)

      _,g,_ = cv2.split(im_output)
      
      if np.mean(g)>100:
        break

    image                      = im_output


    if self.train:
      segmentation_location      = self.segmentation_directory + '/' + self.segmentation_list[index]
      mask_needed                = cv2.imread(segmentation_location,0)
      mask                       = mask_needed.astype(np.float)
      mask                       = mask/255
      mask                       = mask.reshape(512,512,1)

      if self.augmentation:
        ### The augmentations will be done via albumentations library. It always returns a dictionary having 'image' and 'mask' as keys
        sample                   = self.augmentation(image = image, mask = mask)
        image,mask               = sample['image'],sample['mask']

      if self.preprocessing:
        sample                   = self.preprocessing(image = image, mask = mask)
        image,mask               = sample['image'],sample['mask']
      
      return image,mask
    else:
      
      mask                     = np.zeros((512,512,1))
      sample                   = self.preprocessing(image = image,mask=mask)
      image,mask               = sample['image'],sample['mask']
      return image
  
  def __len__(self):
        return len(self.ids)

data4viz = TreeSegmentationDataset(image_directory         = x_train_dir,  
                                    segmentation_directory = y_train_dir
                                  )
image,mask = data4viz[3716]

def visualize(**images):
    """PLot images in one row."""
    n = len(images)
    plt.figure(figsize=(16, 5))
    for i, (name, image) in enumerate(images.items()):
        plt.subplot(1, n, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.title(' '.join(name.split('_')).title())
        plt.imshow(image)
    plt.show()

visualize(trees=image, vegetation_mask=mask.squeeze())

def get_training_augmentations():
  train_transform = [

        albu.HorizontalFlip(p=0.5),

        albu.ShiftScaleRotate(scale_limit=0.5, rotate_limit=0, shift_limit=0.1, p=1, border_mode=0),

        
        albu.IAAAdditiveGaussianNoise(p=0.2),
        albu.IAAPerspective(p=0.5),
        albu.CLAHE(p=1),

        albu.OneOf(
            [

                albu.RandomBrightness(p=1),
#                 albu.RandomGamma(p=1),
            ],
            p=0.9,
        ),

        albu.OneOf(
            [
                albu.IAASharpen(p=1),
                albu.Blur(blur_limit=3, p=1),
                albu.MotionBlur(blur_limit=3, p=1),
            ],
            p=0.9,
        ),

        albu.OneOf(
            [
                albu.RandomContrast(p=1),
                albu.HueSaturationValue(p=1),
            ],
            p=0.9,
        ),
    ]
  return albu.Compose(train_transform)

def get_preprocessing(preprocessing_fn):
    """Construct preprocessing transform
    
    Args:
        preprocessing_fn (callbale): data normalization function 
            (can be specific for each pretrained neural network)
    Return:
        transform: albumentations.Compose
    
    """
    
    _transform = [
        albu.Lambda(image=preprocessing_fn),
        
        albu.Lambda(image=to_tensor, mask=to_tensor),
    ]
    return albu.Compose(_transform)

def to_tensor(x, **kwargs):
    return x.transpose(2, 0, 1).astype('float32')

def get_validation_augmentation():
    """Add paddings to make image shape divisible by 32"""
    test_transform = [
        albu.PadIfNeeded(512, 512)
    ]
    return albu.Compose(test_transform)

# x = get_training_augmentations()
# sample = x(image=image,mask = mask)
# visualize(trees=sample['image'], vegetation_mask=sample['mask'].squeeze())

# visualize(trees = image,vegetation_mask = mask.squeeze())

# augmented_train_dataset = TreeSegmentationDataset(image_directory         = x_train_dir,  
#                                                   segmentation_directory  = y_train_dir,
#                                                   augmentation            = get_training_augmentations()
#                                                   )

# for i in range(3):
#   image, mask = augmented_train_dataset[1]
#   visualize(trees = image,vegetation_mask = mask.squeeze())

ENCODER         = 'efficientnet-b4'
ENCODER_WEIGHTS = 'imagenet'

ACTIVATION      = 'sigmoid' 
DEVICE          = 'cuda'

# create segmentation model with pretrained encoder
model = smp.Unet(
    encoder_name=ENCODER, 
    encoder_weights=ENCODER_WEIGHTS, 
    classes=1, 
    activation=ACTIVATION,
    # decoder_dropout = 0.2,
    
    
)

preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)

train_dataset = TreeSegmentationDataset(
    x_train_dir, 
    y_train_dir, 
    # augmentation=get_training_augmentations(), 
    preprocessing=get_preprocessing(preprocessing_fn),
    
)
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=12)

valid_dataset = TreeSegmentationDataset(
    x_valid_dir, 
    y_valid_dir, 
    # augmentation=get_validation_augmentation(), 
    preprocessing=get_preprocessing(preprocessing_fn),
    
)
valid_loader = DataLoader(valid_dataset, batch_size=16, shuffle=False, num_workers=4)

loss =smp.utils.losses.DiceLoss(eps=1.)

metrics = [
    smp.utils.metrics.IoU(threshold=0.5,eps=1., activation = None),
    smp.utils.metrics.Fscore(threshold=0.5,eps=1., activation = None)
]

optimizer = torch.optim.AdamW([ 
    dict(params=model.parameters(), lr=0.006),
])

train_epoch = smp.utils.train.TrainEpoch(
    model, 
    loss=loss, 
    metrics=metrics, 
    optimizer=optimizer,

    device=DEVICE,
    verbose=True,
)

valid_epoch = smp.utils.train.ValidEpoch(
    model, 
    loss=loss, 
    metrics=metrics, 
    device=DEVICE,
    verbose=True,
)

max_score = 0
extent    = 3
metric2tune = []
count = 0 
for i in range(0, 35):
    
    print('\nEpoch: {}'.format(i))
    train_logs = train_epoch.run(train_loader)
    valid_logs = valid_epoch.run(valid_loader)
    
    # do something (save model, change lr, etc.)
    if max_score < valid_logs['fscore']:
        max_score = valid_logs['fscore']
        torch.save(model, './best_model.pth')
        print('Model saved!')
    
    
    if valid_logs['fscore']>=0.95:
      count = count+1
      if count ==1:
        optimizer.param_groups[0]['lr'] = 0.0001
        print('lr tuned')
      if count%10==0:
        if valid_logs['fscore']<0.955:
          optimizer.param_groups[0]['lr'] = 0.00001
          print('lr tuned')
#         updated_lr,pct_change = lower_learning_rate(metric2tune,extent,optimizer.param_groups[0]['lr'])
#         print(f'Due to pct change being {pct_change},the updated lr is {updated_lr}')
#         optimizer.param_groups[0]['lr'] = updated_lr
    # if i == 25:
    #     optimizer.param_groups[0]['lr'] = 1e-5
    #     print('Decrease decoder learning rate to 1e-5!')

best_model = torch.load('./best_model.pth')

logs = valid_epoch.run(valid_loader)

# test dataset without transformations for image visualization


test_dataset_vis = TreeSegmentationDataset(
    x_valid_dir, 
    y_valid_dir, 
    augmentation=None,
    preprocessing=None
    
)

for i in range(5):
    n = np.random.choice(len(valid_dataset))
    
    image_vis = test_dataset_vis[n][0].astype('uint8')
    image, gt_mask = valid_dataset[n]
    
    gt_mask = gt_mask.squeeze()
    
    x_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0)
    pr_mask = best_model.predict(x_tensor)
    pr_mask = (pr_mask.squeeze().cpu().numpy().round())
        
    visualize(
        image=image_vis, 
        ground_truth_mask=gt_mask, 
        predicted_mask=pr_mask
    )

# pr_mask.shape,gt_mask.shape

test_dataset = TreeSegmentationDataset(x_test_dir,train=False,preprocessing = get_preprocessing(preprocessing_fn))
test_loader = DataLoader(test_dataset, batch_size=1, num_workers=1, shuffle=False, drop_last=False)

predictions = []
for i in tqdm(range(len(test_dataset))):    
    image = test_dataset[i]
    x_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0)
    pr_mask = best_model.predict(x_tensor)
    # print(pr_mask.shape)
    pr_mask = (pr_mask.squeeze().cpu().numpy().round())
    # print(pr_mask.shape)
    predictions.append(pr_mask)    
    # print(i)

# # Change the channel ordering
# image = np.moveaxis(images[0].numpy()/255, 0, -1)

# # Showing the image
# plt.figure()
# plt.subplot(1,2,1)
# plt.imshow(np.moveaxis(image,-1,0), 'gray', interpolation='none')
# plt.subplot(1,2,2)
# plt.imshow(np.moveaxis(image,-1,0), 'gray', interpolation='none')
# plt.imshow(pr_mask, 'jet', interpolation='none', alpha=0.7)
# plt.show()

predictions[0].shape

!rm -rf segmentation
!mkdir segmentation

for n, img in tqdm(enumerate(predictions)):

  img = img*255
  

  _, img = cv2.threshold(img, 128, 255, cv2.THRESH_BINARY)
  # print(img.shape)
  # Making sure the pixels are only 0 and 255 in the image. 
  img = Image.fromarray(img.astype(np.uint8))

  
  img.save(os.path.join("segmentation", f"{n}.png"))

# !pip install aicrowd-cli

%load_ext aicrowd.magic

# %aicrowd login