first commit

app.py

@@ -0,0 +1,275 @@
+import streamlit as st
+import os
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+import torch
+import albumentations as albu
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset as BaseDataset
+from catalyst.runners import SupervisedRunner
+import segmentation_models_pytorch as smp
+from io import StringIO
+
+# streamlit run c:/Users/ronni/Downloads/polyp_seg_web_app/app.py
+
+
+x_test_dir = 'test/test/images'
+y_test_dir = 'test/test/masks'
+ENCODER = 'mobilenet_v2'
+ENCODER_WEIGHTS = 'imagenet'
+CLASSES = ['polyp', 'background']
+ACTIVATION = 'sigmoid'
+
+preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)
+
+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.savefig('x',dpi=400)
+    st.image('x.png')
+
+
+def get_training_augmentation():
+    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.Resize(576, 736, always_apply=True, p=1),
+
+        albu.IAAAdditiveGaussianNoise(p=0.2),
+        albu.IAAPerspective(p=0.5),
+
+        albu.OneOf(
+            [
+                albu.CLAHE(p=1),
+                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_validation_augmentation():
+    """Add paddings to make image shape divisible by 32"""
+    test_transform = [
+        albu.Resize(576, 736)
+    ]
+    return albu.Compose(test_transform)
+
+
+def to_tensor(x, **kwargs):
+    return x.transpose(2, 0, 1).astype('float32')
+
+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)
+
+class Dataset(BaseDataset):
+    """Args:
+        images_dir (str): path to images folder
+        masks_dir (str): path to segmentation masks folder
+        class_values (list): values of classes to extract from segmentation mask
+        augmentation (albumentations.Compose): data transfromation pipeline
+            (e.g. flip, scale, etc.)
+        preprocessing (albumentations.Compose): data preprocessing
+            (e.g. noralization, shape manipulation, etc.)
+
+    """
+
+    CLASSES = ['polyp', 'background']
+
+    def __init__(
+            self,
+            images_dir,
+            masks_dir,
+            classes=None,
+            augmentation=None,
+            preprocessing=None,
+            single_file=False
+    ):
+        
+        if single_file:
+            self.ids = images_dir
+            self.images_fps = os.path.join('test/test/images', self.ids)
+            self.masks_fps = os.path.join('test/test/masks', self.ids)
+        else:
+            self.ids = os.listdir(images_dir)
+            self.images_fps = [os.path.join(images_dir, image_id) for image_id in self.ids]
+            self.masks_fps = [os.path.join(masks_dir, image_id) for image_id in self.ids]
+
+        # convert str names to class values on masks
+        self.class_values = [self.CLASSES.index(cls.lower()) for cls in classes]
+
+        self.augmentation = augmentation
+        self.preprocessing = preprocessing
+
+    def __getitem__(self, i):
+
+        # read data
+        image = cv2.imread(self.images_fps)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        mask = cv2.imread(self.masks_fps, 0)
+        mask[np.where(mask < 8)] = 0
+        mask[np.where(mask > 8)] = 255
+        # extract certain classes from mask (e.g. polyp)
+        masks = [(mask == v) for v in self.class_values]
+        mask = np.stack(masks, axis=-1).astype('float')
+
+        # apply augmentations
+        if self.augmentation:
+            sample = self.augmentation(image=image, mask=mask)
+            image, mask = sample['image'], sample['mask']
+
+        # apply preprocessing
+        if self.preprocessing:
+            sample = self.preprocessing(image=image, mask=mask)
+            image, mask = sample['image'], sample['mask']
+
+        return image, mask
+
+    def __len__(self):
+        return len(self.ids)
+
+def model_infer(img_name):
+    
+    model = smp.UnetPlusPlus(
+        encoder_name=ENCODER, 
+        encoder_weights=ENCODER_WEIGHTS, 
+        encoder_depth=5,
+        decoder_channels=(256, 128, 64, 32, 16),
+        classes=len(CLASSES), 
+        activation=ACTIVATION,
+        decoder_attention_type=None,
+    )
+
+
+    model.load_state_dict(torch.load('best.pth', map_location=torch.device('cpu'))['model_state_dict'])
+    model.eval()
+
+    test_dataset = Dataset(
+        img_name,
+        img_name,
+        augmentation=get_validation_augmentation(),
+        preprocessing=get_preprocessing(preprocessing_fn),
+        classes=CLASSES,
+        single_file=True
+    )
+
+    test_dataloader = DataLoader(test_dataset)
+
+    loaders = {"infer": test_dataloader}
+
+    runner = SupervisedRunner()
+
+    logits = []
+    f = 0
+    for prediction in runner.predict_loader(model=model, loader=loaders['infer'],cpu=True):
+        if f < 3:
+            logits.append(prediction['logits'])
+            f = f + 1
+        else:
+            break
+
+    threshold = 0.5
+    break_at = 1
+
+    for i, (input, output) in enumerate(zip(
+            test_dataset, logits)):
+        image, mask = input
+
+        image_vis = image.transpose(1, 2, 0)
+        gt_mask = mask[0].astype('uint8')
+        pr_mask = (output[0].numpy() > threshold).astype('uint8')[0]
+        i = i + 1
+        if i >= break_at:
+            break
+    
+    return image_vis, gt_mask, pr_mask
+PAGE_TITLE = "Polyp Segmentation" 
+
+def file_selector(folder_path='.'):
+    filenames = os.listdir(folder_path)
+    selected_filename = st.selectbox('Select a file', filenames)
+    return os.path.join(folder_path, selected_filename)
+
+def file_selector_ui():
+    folder_path = './test/test/images'
+    filename = file_selector(folder_path=folder_path)
+    printname = list(filename)
+    printname[filename.rfind('\\')] = '/'
+    st.write('You selected`%s`' % ''.join(printname))
+    return filename
+
+def file_upload(folder_path='.'):
+    filenames = os.listdir(folder_path)
+    folder_path = './test/test/images'
+    uploaded_file = st.file_uploader("Choose a file")
+    filename =     os.path.join(folder_path, uploaded_file.name)
+    printname = list(filename)
+    printname[filename.rfind('\\')] = '/'
+    st.write('You selected`%s`' % ''.join(printname))
+    return filename
+
+
+def main():
+    st.set_page_config(page_title=PAGE_TITLE, layout="wide")
+    st.title(PAGE_TITLE)
+    image_path = file_selector_ui()
+#    image_path = file_upload()
+    image_path = os.path.abspath(image_path)
+    to_infer = image_path[image_path.rfind("\\") + 1:]
+
+    if os.path.isfile(image_path) is True:
+        _, file_extension = os.path.splitext(image_path)
+        if file_extension == ".jpg":
+            image_vis, gt_mask, pr_mask = model_infer(to_infer)
+            visualize(
+                image=image_vis, 
+                ground_truth_mask=gt_mask, 
+                predicted_mask=pr_mask
+            )            
+
+if __name__ == "__main__":
+    main()

requirements.txt.txt

@@ -0,0 +1,14 @@
+catalyst==19.04rc1
+git+https://github.com/albu/albumentations@bdd6a4e
+git+https://github.com/qubvel/segmentation_models.pytorch
+os
+numpy
+cv2
+matplotlib
+torch
+albumentations
+segmentation_models_pytorch
+collections
+splitfolders
+gc
+streamlit