Add files via upload

Author: DebeshJha

crf.py

@@ -0,0 +1,44 @@
+
+import os
+import numpy as np
+import cv2
+import pydensecrf.densecrf as dcrf
+from pydensecrf.utils import unary_from_labels, create_pairwise_bilateral
+
+def apply_crf(ori_image, mask):
+    """ Conditional Random Field
+    ori_image: np.array with value between 0-255
+    mask: np.array with value between 0-1
+    """
+
+    ## Grayscale to RGB
+    # if len(mask.shape) < 3:
+    #     mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
+
+    ## Converting the anotations RGB to single 32  bit color
+    annotated_label = mask.astype(np.int32)
+    # annotated_label = mask[:,:,0] + (mask[:,:,1]<<8) + (mask[:,:,2]<<16)
+
+    ## Convert the 32bit integer color to 0,1, 2, ... labels.
+    colors, labels = np.unique(annotated_label, return_inverse=True)
+    n_labels = 2
+
+    ## Setting up the CRF model
+    d = dcrf.DenseCRF2D(ori_image.shape[1], ori_image.shape[0], n_labels)
+
+    ## Get unary potentials (neg log probability)
+    U = unary_from_labels(labels, n_labels, gt_prob=0.7, zero_unsure=False)
+    d.setUnaryEnergy(U)
+
+    ## This adds the color-independent term, features are the locations only.
+    d.addPairwiseGaussian(sxy=(3, 3), compat=3, kernel=dcrf.DIAG_KERNEL, normalization=dcrf.NORMALIZE_SYMMETRIC)
+
+    ## Run Inference for 10 steps
+    Q = d.inference(10)
+
+    ## Find out the most probable class for each pixel.
+    MAP = np.argmax(Q, axis=0)
+
+    return MAP.reshape((ori_image.shape[0], ori_image.shape[1]))
+
+

data.py

@@ -0,0 +1,78 @@
+
+import os
+import numpy as np
+import cv2
+from glob import glob
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+def load_names(path, file_path):
+    f = open(file_path, "r")
+    data = f.read().split("\n")[:-1]
+    images = [os.path.join(path,"images", name) + ".jpg" for name in data]
+    masks = [os.path.join(path,"masks", name) + ".jpg" for name in data]
+    return images, masks
+
+def load_data(path):
+    train_names_path = f"{path}/train.txt"
+    valid_names_path = f"{path}/val.txt"
+
+    train_x, train_y = load_names(path, train_names_path)
+    valid_x, valid_y = load_names(path, valid_names_path)
+
+    return (train_x, train_y), (valid_x, valid_y)
+
+class KvasirDataset(Dataset):
+    """ Dataset for the Kvasir-SEG dataset. """
+    def __init__(self, images_path, masks_path, size):
+        """
+        Arguments:
+            images_path: A list of path of the images.
+            masks_path: A list of path of the masks.
+        """
+
+        self.images_path = images_path
+        self.masks_path = masks_path
+        self.height = size[0]
+        self.width = size[1]
+        self.n_samples = len(images_path)
+
+    def __getitem__(self, index):
+        """ Reading image and mask. """
+        image = cv2.imread(self.images_path[index], cv2.IMREAD_COLOR)
+        mask = cv2.imread(self.masks_path[index], cv2.IMREAD_GRAYSCALE)
+
+        """ Resizing. """
+        image1 = cv2.resize(image, (self.width, self.height))
+        # image2 = cv2.resize(image, (self.width//2, self.height//2))
+        # image3 = cv2.resize(image, (self.width//4, self.height//4))
+        mask = cv2.resize(mask, (self.width, self.height))
+
+        """ Proper channel formatting. """
+        image1 = np.transpose(image1, (2, 0, 1))
+        # image2 = np.transpose(image2, (2, 0, 1))
+        # image3 = np.transpose(image3, (2, 0, 1))
+        mask = np.expand_dims(mask, axis=0)
+
+        """ Normalization. """
+        image1 = image1/255.0
+        # image2 = image2/255.0
+        # image3 = image3/255.0
+        mask = mask/255.0
+
+        """ Changing datatype to float32. """
+        image1 = image1.astype(np.float32)
+        # image2 = image2.astype(np.float32)
+        # image3 = image3.astype(np.float32)
+        mask = mask.astype(np.float32)
+
+        """ Changing numpy to tensor. """
+        image1 = torch.from_numpy(image1)
+        # image2 = torch.from_numpy(image2)
+        # image3 = torch.from_numpy(image3)
+        mask = torch.from_numpy(mask)
+
+        return image1, mask
+
+    def __len__(self):
+        return self.n_samples

data_aug.py

@@ -0,0 +1,177 @@
+
+import os
+import random
+import numpy as np
+import cv2
+from tqdm import tqdm
+from glob import glob
+from sklearn.model_selection import train_test_split
+from utils import create_dir
+from data import load_data
+
+from albumentations import (
+    PadIfNeeded,
+    HorizontalFlip,
+    VerticalFlip,
+    CenterCrop,
+    Crop,
+    RandomCrop,
+    Compose,
+    Transpose,
+    RandomRotate90,
+    ElasticTransform,
+    GridDistortion,
+    OpticalDistortion,
+    RandomSizedCrop,
+    OneOf,
+    CLAHE,
+    RandomBrightnessContrast,
+    RandomGamma,
+    HueSaturationValue,
+    RGBShift,
+    RandomBrightness,
+    RandomContrast,
+    MotionBlur,
+    MedianBlur,
+    GaussianBlur,
+    GaussNoise,
+    ChannelShuffle,
+    CoarseDropout
+)
+
+def augment_data(images, masks, save_path, augment=True):
+    """ Performing data augmentation. """
+    size = (512, 512)
+    crop_size = (448, 448)
+
+    for idx, (x, y) in tqdm(enumerate(zip(images, masks)), total=len(images)):
+        image_name = x.split("/")[-1].split(".")[0]
+        mask_name = y.split("/")[-1].split(".")[0]
+
+        x = cv2.imread(x, cv2.IMREAD_COLOR)
+        y = cv2.imread(y, cv2.IMREAD_COLOR)
+
+        if x.shape[0] >= size[0] and x.shape[1] >= size[1]:
+            if augment == True:
+                ## Crop
+                x_min = 0
+                y_min = 0
+                x_max = x_min + size[0]
+                y_max = y_min + size[1]
+
+                aug = Crop(p=1, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max)
+                augmented = aug(image=x, mask=y)
+                x1 = augmented['image']
+                y1 = augmented['mask']
+
+                # Random Rotate 90 degree
+                aug = RandomRotate90(p=1)
+                augmented = aug(image=x, mask=y)
+                x2 = augmented['image']
+                y2 = augmented['mask']
+
+                ## ElasticTransform
+                aug = ElasticTransform(p=1, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03)
+                augmented = aug(image=x, mask=y)
+                x3 = augmented['image']
+                y3 = augmented['mask']
+
+                ## Grid Distortion
+                aug = GridDistortion(p=1)
+                augmented = aug(image=x, mask=y)
+                x4 = augmented['image']
+                y4 = augmented['mask']
+
+                ## Optical Distortion
+                aug = OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5)
+                augmented = aug(image=x, mask=y)
+                x5 = augmented['image']
+                y5 = augmented['mask']
+
+                ## Vertical Flip
+                aug = VerticalFlip(p=1)
+                augmented = aug(image=x, mask=y)
+                x6 = augmented['image']
+                y6 = augmented['mask']
+
+                ## Horizontal Flip
+                aug = HorizontalFlip(p=1)
+                augmented = aug(image=x, mask=y)
+                x7 = augmented['image']
+                y7 = augmented['mask']
+
+                ## Grayscale
+                x8 = cv2.cvtColor(x, cv2.COLOR_RGB2GRAY)
+                y8 = y
+
+                aug = RGBShift(p=1)
+                augmented = aug(image=x, mask=y)
+                x9 = augmented['image']
+                y9 = augmented['mask']
+
+                aug = ChannelShuffle(p=1)
+                augmented = aug(image=x, mask=y)
+                x10 = augmented['image']
+                y10 = augmented['mask']
+
+                aug = CoarseDropout(p=1, max_holes=10, max_height=32, max_width=32)
+                augmented = aug(image=x, mask=y)
+                x11 = augmented['image']
+                y11 = augmented['mask']
+
+                aug = GaussNoise(p=1)
+                augmented = aug(image=x, mask=y)
+                x12 = augmented['image']
+                y12 = augmented['mask']
+
+                images = [
+                    x, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12
+                ]
+                masks  = [
+                    y, y1, y2, y3, y4, y5, y6, y7, y8, y9, y10, y11, y12
+                ]
+
+            else:
+                images = [x]
+                masks  = [y]
+
+            idx = 0
+        for i, m in zip(images, masks):
+            i = cv2.resize(i, size)
+            m = cv2.resize(m, size)
+
+            if len(images) == 1:
+                tmp_image_name = f"{image_name}.jpg"
+                tmp_mask_name  = f"{mask_name}.jpg"
+            else:
+                tmp_image_name = f"{image_name}_{idx}.jpg"
+                tmp_mask_name  = f"{mask_name}_{idx}.jpg"
+
+            image_path = os.path.join(save_path, "image/", tmp_image_name)
+            mask_path  = os.path.join(save_path, "mask/", tmp_mask_name)
+
+            cv2.imwrite(image_path, i)
+            cv2.imwrite(mask_path, m)
+
+            idx += 1
+
+def main():
+    np.random.seed(42)
+
+    path = "/home/nikhilroxtomar/lab/DATA/Kvasir-SEG/"
+    # path = "/media/nikhil/ML/ml_dataset/Kvasir-SEG/"
+    (train_x, train_y), (test_x, test_y) = load_data(path)
+
+    print("Train: ", len(train_x))
+    print("Valid: ", len(test_x))
+
+    create_dir("new_data/train/image/")
+    create_dir("new_data/train/mask/")
+    create_dir("new_data/test/image/")
+    create_dir("new_data/test/mask/")
+
+    augment_data(train_x, train_y, "new_data/train/", augment=False)
+    augment_data(test_x, test_y, "new_data/test/", augment=False)
+
+if __name__ == "__main__":
+    main()

loss.py

@@ -0,0 +1,90 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class DiceLoss(nn.Module):
+    def __init__(self, weight=None, size_average=True):
+        super(DiceLoss, self).__init__()
+
+    def forward(self, inputs, targets, smooth=1):
+
+        #comment out if your model contains a sigmoid or equivalent activation layer
+        inputs = torch.sigmoid(inputs)
+
+        #flatten label and prediction tensors
+        inputs = inputs.view(-1)
+        targets = targets.view(-1)
+
+        intersection = (inputs * targets).sum()
+        dice = (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)
+
+        return 1 - dice
+
+class DiceBCELoss(nn.Module):
+    def __init__(self, weight=None, size_average=True):
+        super(DiceBCELoss, self).__init__()
+
+    def forward(self, inputs, targets, smooth=1):
+
+        #comment out if your model contains a sigmoid or equivalent activation layer
+        inputs = torch.sigmoid(inputs)
+
+        #flatten label and prediction tensors
+        inputs = inputs.view(-1)
+        targets = targets.view(-1)
+
+        intersection = (inputs * targets).sum()
+        dice_loss = 1 - (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)
+        BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')
+        Dice_BCE = BCE + dice_loss
+
+        return Dice_BCE
+
+class IoULoss(nn.Module):
+    def __init__(self, weight=None, size_average=True):
+        super(IoULoss, self).__init__()
+
+    def forward(self, inputs, targets, smooth=1):
+
+        #comment out if your model contains a sigmoid or equivalent activation layer
+        inputs = torch.sigmoid(inputs)
+
+        #flatten label and prediction tensors
+        inputs = inputs.view(-1)
+        targets = targets.view(-1)
+
+        #intersection is equivalent to True Positive count
+        #union is the mutually inclusive area of all labels & predictions
+        intersection = (inputs * targets).sum()
+        total = (inputs + targets).sum()
+        union = total - intersection
+
+        IoU = (intersection + smooth)/(union + smooth)
+
+        return -IoU
+
+class IoUBCELoss(nn.Module):
+    def __init__(self, weight=None, size_average=True):
+        super(IoUBCELoss, self).__init__()
+
+    def forward(self, inputs, targets, smooth=1):
+
+        #comment out if your model contains a sigmoid or equivalent activation layer
+        inputs = torch.sigmoid(inputs)
+
+        #flatten label and prediction tensors
+        inputs = inputs.view(-1)
+        targets = targets.view(-1)
+
+        #intersection is equivalent to True Positive count
+        #union is the mutually inclusive area of all labels & predictions
+        intersection = (inputs * targets).sum()
+        total = (inputs + targets).sum()
+        union = total - intersection
+
+        IoU = - (intersection + smooth)/(union + smooth)
+
+        BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')
+        IoU_BCE = BCE + IoU
+
+        return IoU_BCE