├── mosaic ├── __init__.py ├── math_fit.py ├── create_mean_images.py ├── mean_hist.py ├── hist_match.py ├── mosaic_api_examples.py.ipynb ├── create_mosaic.py └── mosaic_api.py ├── src ├── __init__.py ├── nick_zoo │ ├── __init__.py │ ├── utils.py │ ├── resnet_blocks.py │ ├── unet_flex_fpn.py │ ├── layers.py │ └── unet_flex.py ├── romul_zoo │ └── __init__.py ├── config.py ├── make_folds.py ├── utils.py ├── lr_scheduler.py ├── dataset.py ├── metrics.py ├── losses.py ├── argus_models.py ├── nn_modules.py ├── lovasz.py └── transforms.py ├── .dockerignore ├── data ├── README.md ├── mean_hist.npy └── .gitignore ├── readme_img ├── architecture.png └── postprocess.png ├── requirements.txt ├── Dockerfile ├── run_pipeline.sh ├── Makefile ├── pipeline ├── postprocessed_submission.py ├── mean_submission.py ├── mos-fpn-lovasz-se-resnext50-001 │ ├── train_folds.py │ └── predict_folds.py ├── fpn-lovasz-se-resnext50-006 │ ├── train_folds.py │ └── predict_folds.py └── fpn-lovasz-se-resnext50-006-after-001 │ ├── train_folds.py │ └── predict_folds.py ├── LICENSE ├── .gitignore ├── make_inpaint_images.py ├── unused ├── make_mosaic_inpaint_images.py ├── make_double_inpaint_images.py └── train_mean_teacher_folds.py ├── train_folds.py ├── random_search.py ├── after_train_folds.py ├── notebooks ├── cv_check.ipynb ├── mean_submission.ipynb ├── make_submission.ipynb ├── log_visualize.ipynb └── train.ipynb ├── README.md └── predict_folds.py /mosaic/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.dockerignore: -------------------------------------------------------------------------------- 1 | data 2 | -------------------------------------------------------------------------------- /data/README.md: -------------------------------------------------------------------------------- 1 | Data folder -------------------------------------------------------------------------------- /src/nick_zoo/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/romul_zoo/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data/mean_hist.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lRomul/argus-tgs-salt/HEAD/data/mean_hist.npy -------------------------------------------------------------------------------- /readme_img/architecture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lRomul/argus-tgs-salt/HEAD/readme_img/architecture.png -------------------------------------------------------------------------------- /readme_img/postprocess.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lRomul/argus-tgs-salt/HEAD/readme_img/postprocess.png -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | pycocotools==2.0.0 2 | torchsummary 3 | scikit-optimize 4 | pretrainedmodels 5 | pytorch-argus==0.0.5 6 | cffi 7 | tqdm 8 | shapely 9 | -------------------------------------------------------------------------------- /data/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | !train_folds.csv 3 | !mean_hist.npy 4 | !train_folds_148.csv 5 | !mosaic/ 6 | !mosaic/pazzles_6013.csv 7 | !mosaic/id2pred_fold.json 8 | !train_folds_148_mos_emb_1.csv 9 | !.gitignore 10 | !README.md 11 | -------------------------------------------------------------------------------- /src/nick_zoo/utils.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | 3 | def class_for_name(module_name, class_name): 4 | # load the module, will raise ImportError if module cannot be loaded 5 | m = importlib.import_module(module_name) 6 | # get the class, will raise AttributeError if class cannot be found 7 | return getattr(m, class_name) 8 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM floydhub/pytorch:0.4.1-gpu.cuda9cudnn7-py3.34 2 | 3 | # Install python packages 4 | COPY requirements.txt requirements.txt 5 | RUN pip3 install --upgrade pip 6 | RUN pip install --no-cache-dir -r requirements.txt 7 | 8 | ENV PYTHONPATH $PYTHONPATH:/workdir 9 | ENV TORCH_HOME=/workdir/data/.torch 10 | 11 | WORKDIR /workdir 12 | -------------------------------------------------------------------------------- /src/config.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | 3 | DATA_DIR = '/workdir/data/' 4 | TRAIN_DIR = join(DATA_DIR, 'train') 5 | TEST_DIR = join(DATA_DIR, 'test') 6 | DEPTHS_PATH = join(DATA_DIR, 'depths.csv') 7 | TRAIN_FOLDS_PATH = join(DATA_DIR, 'train_folds.csv') 8 | TRAIN_CSV_PATH = join(DATA_DIR, 'train.csv') 9 | SAMPLE_SUBM_PATH = join(DATA_DIR, 'sample_submission.csv') 10 | MEAN_HIST_PATH = join(DATA_DIR, 'mean_hist.npy') 11 | N_FOLDS = 6 12 | -------------------------------------------------------------------------------- /mosaic/math_fit.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | def fit_third_order(x1, x2, y1, y2, dy1, dy2): 5 | if x2 != x1: 6 | a = (x1*dy1 + x1*dy2 - x2*dy1 - x2*dy2 - 2*y1 + 2*y2)/(x1 - x2)**3 7 | b = (x1**2*(-dy1) - 2*x1**2*dy2 - x1*x2*dy1 + x1*x2*dy2 + 3*x1*y1 - 3*x1*y2 + 2*x2**2*dy1 + x2**2*dy2 + 3*x2*y1 - 3*x2*y2)/(x1 - x2)**3 8 | c = (x1**3*dy2 + 2*x1**2*x2*dy1 + x1**2*x2*dy2 - x1*x2**2*dy1 - 2*x1*x2**2*dy2 - 6*x1*x2*y1 + 6*x1*x2*y2 - x2**3*dy1)/(x1 - x2)**3 9 | d = (x1**3*(-x2)*dy2 + x1**3*y2 - x1**2*x2**2*dy1 + x1**2*x2**2*dy2 - 3*x1**2*x2*y2 + x1*x2**3*dy1 + 3*x1*x2**2*y1 - x2**3*y1)/(x1 - x2)**3 10 | return a, b, c, d 11 | return None 12 | -------------------------------------------------------------------------------- /run_pipeline.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | set -e 3 | 4 | # make biharmonic inpaint 5 | python make_inpaint_images.py 6 | 7 | # train fpn-lovasz-se-resnext50-006-after-001 8 | python pipeline/fpn-lovasz-se-resnext50-006/train_folds.py 9 | python pipeline/fpn-lovasz-se-resnext50-006-after-001/train_folds.py 10 | python pipeline/fpn-lovasz-se-resnext50-006-after-001/predict_folds.py 11 | 12 | # train mos-fpn-lovasz-se-resnext50-001 13 | python pipeline/mos-fpn-lovasz-se-resnext50-001/train_folds.py 14 | python pipeline/mos-fpn-lovasz-se-resnext50-001/predict_folds.py 15 | 16 | # make mean submission 17 | python pipeline/mean_submission.py 18 | 19 | # make postprocessed submission 20 | python pipeline/postprocessed_submission.py 21 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | NAME?=argus-tgs-salt 2 | COMMAND?=bash 3 | 4 | GPUS?=all 5 | ifeq ($(GPUS),none) 6 | GPUS_OPTION= 7 | else 8 | GPUS_OPTION=--gpus=$(GPUS) 9 | endif 10 | 11 | .PHONY: all 12 | all: stop build run 13 | 14 | .PHONY: build 15 | build: 16 | docker build -t $(NAME) . 17 | 18 | .PHONY: stop 19 | stop: 20 | -docker stop $(NAME) 21 | -docker rm $(NAME) 22 | 23 | .PHONY: run 24 | run: 25 | docker run --rm -dit \ 26 | $(GPUS_OPTION) \ 27 | --net=host \ 28 | --ipc=host \ 29 | -v $(shell pwd):/workdir \ 30 | --name=$(NAME) \ 31 | $(NAME) \ 32 | $(COMMAND) 33 | docker attach $(NAME) 34 | 35 | .PHONY: attach 36 | attach: 37 | docker attach $(NAME) 38 | 39 | .PHONY: logs 40 | logs: 41 | docker logs -f $(NAME) 42 | 43 | .PHONY: exec 44 | exec: 45 | docker exec -it $(NAME) $(COMMAND) 46 | -------------------------------------------------------------------------------- /mosaic/create_mean_images.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import cv2 3 | import tqdm 4 | import numpy as np 5 | from os.path import join 6 | from mosaic.hist_match import hist_match 7 | from src.config import MEAN_HIST_PATH, TRAIN_DIR, TEST_DIR 8 | from src.utils import make_dir 9 | 10 | IMAGES_DIR_NAME = 'images_hist_mean' 11 | 12 | 13 | mean_hist = np.load(MEAN_HIST_PATH) 14 | 15 | images_paths = glob.glob(join(TRAIN_DIR, 'images', '*.png')) 16 | images_paths += glob.glob(join(TEST_DIR, 'images', '*.png')) 17 | 18 | make_dir(join(TEST_DIR, IMAGES_DIR_NAME)) 19 | make_dir(join(TRAIN_DIR, IMAGES_DIR_NAME)) 20 | 21 | 22 | for image_path in tqdm.tqdm(images_paths): 23 | image = cv2.imread(image_path) 24 | hist_image = hist_match(image, mean_hist) 25 | hist_image = np.round(hist_image).astype(np.uint8) 26 | save_image_path = image_path.replace('/images/', f'/{IMAGES_DIR_NAME}/') 27 | cv2.imwrite(save_image_path, hist_image) 28 | -------------------------------------------------------------------------------- /pipeline/postprocessed_submission.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import pandas as pd 4 | from os.path import join 5 | 6 | from src import config 7 | from src.utils import RLenc, make_dir 8 | from pipeline.mean_submission import MEAN_PREDICTION_DIR 9 | from mosaic.postprocess import postprocess 10 | 11 | mean_path = join(MEAN_PREDICTION_DIR, 'masks') 12 | postprocess_path = join(MEAN_PREDICTION_DIR, 'postprocessed') 13 | 14 | make_dir(postprocess_path) 15 | 16 | 17 | if __name__ == "__main__": 18 | print('Make postprocessed submission') 19 | 20 | postprocess(mean_path, postprocess_path) 21 | 22 | sample_submition = pd.read_csv(config.SAMPLE_SUBM_PATH) 23 | for i, row in sample_submition.iterrows(): 24 | pred_name = row.id+'.png' 25 | pred = cv2.imread(join(postprocess_path, pred_name), 26 | cv2.IMREAD_GRAYSCALE) 27 | pred = pred > 0 28 | rle_mask = RLenc(pred.astype(np.uint8)) 29 | row.rle_mask = rle_mask 30 | 31 | sample_submition.to_csv(join(MEAN_PREDICTION_DIR, 32 | 'submission.csv'), index=False) 33 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Ruslan Baikulov, Nikolay Falaleev 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /mosaic/mean_hist.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import tqdm 4 | import numpy as np 5 | 6 | from src.config import MEAN_HIST_PATH, TRAIN_DIR, TEST_DIR 7 | 8 | img_train_dir = os.path.join(TRAIN_DIR, 'images') 9 | img_test_dir = os.path.join(TEST_DIR, 'images') 10 | 11 | images_names = os.listdir(img_train_dir) 12 | images_names += os.listdir(img_test_dir) 13 | 14 | n_samples = len(images_names) 15 | 16 | images_paths = [] 17 | for img_name in images_names: 18 | train_path = os.path.join(img_train_dir, img_name) 19 | test_path = os.path.join(img_test_dir, img_name) 20 | if os.path.exists(train_path): 21 | images_paths.append(train_path) 22 | elif os.path.exists(test_path): 23 | images_paths.append(test_path) 24 | 25 | # Check if everything is ok 26 | print("Samples:", n_samples, "Ok:", len(images_paths)==len(images_names)) 27 | 28 | imgs = [cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) 29 | for img_path in images_paths] 30 | 31 | hist = np.zeros(256, dtype=np.int64) 32 | 33 | for img in tqdm.tqdm(imgs, desc="Create histograms"): 34 | if np.sum(img): 35 | values, counts = np.unique(img, return_counts=True) 36 | for (v, c) in zip(values, counts): 37 | hist[v] += c 38 | 39 | np.save(MEAN_HIST_PATH, hist/len(images_names)) 40 | print("Mean histogram saved!") 41 | -------------------------------------------------------------------------------- /src/nick_zoo/resnet_blocks.py: -------------------------------------------------------------------------------- 1 | # Resnet family definition dicts for help and as a reference 2 | 3 | resnet18 = [ 4 | {'type': 'basic', 'n_blocks': 2, 'n_filters': 64, 'stride': 1}, 5 | {'type': 'basic', 'n_blocks': 2, 'n_filters': 128, 'stride': 2}, 6 | {'type': 'basic', 'n_blocks': 2, 'n_filters': 256, 'stride': 2}, 7 | {'type': 'basic', 'n_blocks': 2, 'n_filters': 512, 'stride': 2}, 8 | ] 9 | 10 | resnet34 = [ 11 | {'type': 'basic', 'n_blocks': 3, 'n_filters': 64, 'stride': 1}, 12 | {'type': 'basic', 'n_blocks': 4, 'n_filters': 128, 'stride': 2}, 13 | {'type': 'basic', 'n_blocks': 6, 'n_filters': 256, 'stride': 2}, 14 | {'type': 'basic', 'n_blocks': 3, 'n_filters': 512, 'stride': 2}, 15 | ] 16 | 17 | resnet50 = [ 18 | {'type': 'bottleneck', 'n_blocks': 3, 'n_filters': 64, 'stride': 1}, 19 | {'type': 'bottleneck', 'n_blocks': 4, 'n_filters': 128, 'stride': 2}, 20 | {'type': 'bottleneck', 'n_blocks': 6, 'n_filters': 256, 'stride': 2}, 21 | {'type': 'bottleneck', 'n_blocks': 3, 'n_filters': 512, 'stride': 2}, 22 | ] 23 | 24 | resnet101 = [ 25 | {'type': 'bottleneck', 'n_blocks': 3, 'n_filters': 64, 'stride': 1}, 26 | {'type': 'bottleneck', 'n_blocks': 4, 'n_filters': 128, 'stride': 2}, 27 | {'type': 'bottleneck', 'n_blocks': 23, 'n_filters': 256, 'stride': 2}, 28 | {'type': 'bottleneck', 'n_blocks': 3, 'n_filters': 512, 'stride': 2}, 29 | ] 30 | 31 | resnet152 = [ 32 | {'type': 'bottleneck', 'n_blocks': 3, 'n_filters': 64, 'stride': 1}, 33 | {'type': 'bottleneck', 'n_blocks': 8, 'n_filters': 128, 'stride': 2}, 34 | {'type': 'bottleneck', 'n_blocks': 36, 'n_filters': 256, 'stride': 2}, 35 | {'type': 'bottleneck', 'n_blocks': 3, 'n_filters': 512, 'stride': 2}, 36 | ] 37 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | 3 | # Byte-compiled / optimized / DLL files 4 | __pycache__/ 5 | *.py[cod] 6 | *$py.class 7 | 8 | # C extensions 9 | *.so 10 | 11 | # Distribution / packaging 12 | .Python 13 | build/ 14 | develop-eggs/ 15 | dist/ 16 | downloads/ 17 | eggs/ 18 | .eggs/ 19 | lib/ 20 | lib64/ 21 | parts/ 22 | sdist/ 23 | var/ 24 | wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | .hypothesis/ 50 | .pytest_cache/ 51 | 52 | # Translations 53 | *.mo 54 | *.pot 55 | 56 | # Django stuff: 57 | *.log 58 | local_settings.py 59 | db.sqlite3 60 | 61 | # Flask stuff: 62 | instance/ 63 | .webassets-cache 64 | 65 | # Scrapy stuff: 66 | .scrapy 67 | 68 | # Sphinx documentation 69 | docs/_build/ 70 | 71 | # PyBuilder 72 | target/ 73 | 74 | # Jupyter Notebook 75 | .ipynb_checkpoints 76 | 77 | # pyenv 78 | .python-version 79 | 80 | # celery beat schedule file 81 | celerybeat-schedule 82 | 83 | # SageMath parsed files 84 | *.sage.py 85 | 86 | # Environments 87 | .env 88 | .venv 89 | env/ 90 | venv/ 91 | ENV/ 92 | env.bak/ 93 | venv.bak/ 94 | 95 | # Spyder project settings 96 | .spyderproject 97 | .spyproject 98 | 99 | # Rope project settings 100 | .ropeproject 101 | 102 | # mkdocs documentation 103 | /site 104 | 105 | # mypy 106 | .mypy_cache/ 107 | -------------------------------------------------------------------------------- /src/make_folds.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | import pandas as pd 3 | import numpy as np 4 | import random 5 | import cv2 6 | 7 | from src import config 8 | 9 | 10 | def get_correct_train_ids(train_csv_path, train_dir): 11 | train_df = pd.read_csv(train_csv_path, index_col='id') 12 | train_df.fillna('', inplace=True) 13 | correct_ids = [] 14 | 15 | for index, row in train_df.iterrows(): 16 | image_path = join(train_dir, 'images', index + '.png') 17 | image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) 18 | pixel_sum = np.sum(image) 19 | if pixel_sum: 20 | correct_ids.append(index) 21 | return correct_ids 22 | 23 | 24 | def make_train_folds(train_csv_path, train_dir, depths_path, n_folds): 25 | depths_df = pd.read_csv(depths_path, index_col='id') 26 | train_ids = get_correct_train_ids(train_csv_path, train_dir) 27 | depths_df = depths_df.loc[train_ids] 28 | depths_df['image_path'] = depths_df.index.map( 29 | lambda x: join(train_dir, 'images', x + '.png')) 30 | depths_df['mask_path'] = depths_df.index.map( 31 | lambda x: join(train_dir, 'masks', x + '.png')) 32 | depths_df.sort_values('z', inplace=True) 33 | depths_df['fold'] = [i % n_folds for i in range(depths_df.shape[0])] 34 | depths_df.sort_index(inplace=True) 35 | return depths_df 36 | 37 | 38 | if __name__ == '__main__': 39 | random.seed(42) 40 | np.random.seed(42) 41 | 42 | train_folds_df = make_train_folds(config.TRAIN_CSV_PATH, 43 | config.TRAIN_DIR, 44 | config.DEPTHS_PATH, 45 | config.N_FOLDS) 46 | train_folds_df.to_csv(config.TRAIN_FOLDS_PATH) 47 | print(f"Folds saved to {config.TRAIN_FOLDS_PATH}") 48 | -------------------------------------------------------------------------------- /src/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | 4 | 5 | def make_dir(dir_path): 6 | if not os.path.exists(dir_path): 7 | os.makedirs(dir_path) 8 | 9 | 10 | def filename_without_ext(file_path): 11 | basename = os.path.basename(file_path) 12 | filename = os.path.splitext(basename)[0] 13 | return filename 14 | 15 | 16 | # Source https://www.kaggle.com/bguberfain/unet-with-depth 17 | def RLenc(img, order='F', format=True): 18 | """ 19 | img is binary mask image, shape (r,c) 20 | order is down-then-right, i.e. Fortran 21 | format determines if the order needs to be preformatted (according to submission rules) or not 22 | 23 | returns run length as an array or string (if format is True) 24 | """ 25 | bytes = img.reshape(img.shape[0] * img.shape[1], order=order) 26 | runs = [] ## list of run lengths 27 | r = 0 ## the current run length 28 | pos = 1 ## count starts from 1 per WK 29 | for c in bytes: 30 | if (c == 0): 31 | if r != 0: 32 | runs.append((pos, r)) 33 | pos += r 34 | r = 0 35 | pos += 1 36 | else: 37 | r += 1 38 | 39 | # if last run is unsaved (i.e. data ends with 1) 40 | if r != 0: 41 | runs.append((pos, r)) 42 | pos += r 43 | r = 0 44 | 45 | if format: 46 | z = '' 47 | 48 | for rr in runs: 49 | z += '{} {} '.format(rr[0], rr[1]) 50 | return z[:-1] 51 | else: 52 | return runs 53 | 54 | 55 | def sigmoid_rampup(current, rampup_length): 56 | """Exponential rampup from https://arxiv.org/abs/1610.02242""" 57 | if rampup_length == 0: 58 | return 1.0 59 | else: 60 | current = np.clip(current, 0.0, rampup_length) 61 | phase = 1.0 - current / rampup_length 62 | return float(np.exp(-5.0 * phase * phase)) 63 | -------------------------------------------------------------------------------- /mosaic/hist_match.py: -------------------------------------------------------------------------------- 1 | # From https://stackoverflow.com/questions/32655686/histogram-matching-of-two-images-in-python-2-x 2 | 3 | import numpy as np 4 | 5 | from src.config import MEAN_HIST_PATH 6 | 7 | 8 | def hist_match(source, hist): 9 | """ 10 | Adjust the pixel values of a grayscale image such that its histogram 11 | matches a given one 12 | 13 | Arguments: 14 | ----------- 15 | source: np.ndarray 16 | Image to transform; the histogram is computed over the flattened 17 | array 18 | template: np.ndarray 19 | Template histogram 20 | Returns: 21 | ----------- 22 | matched: np.ndarray 23 | The transformed output image 24 | """ 25 | 26 | oldshape = source.shape 27 | source = source.ravel() 28 | 29 | # get the set of unique pixel values and their corresponding indices and 30 | # counts 31 | s_values, bin_idx, s_counts = np.unique(source, return_inverse=True, 32 | return_counts=True) 33 | t_values = np.linspace(0, 255, 256).astype(np.int) 34 | 35 | # take the cumsum of the counts and normalize by the number of pixels to 36 | # get the empirical cumulative distribution functions for the source and 37 | # template images (maps pixel value --> quantile) 38 | s_quantiles = np.cumsum(s_counts).astype(np.float64) 39 | s_quantiles /= s_quantiles[-1] 40 | t_quantiles = np.cumsum(hist).astype(np.float64) 41 | t_quantiles /= t_quantiles[-1] 42 | 43 | # interpolate linearly to find the pixel values in the template image 44 | # that correspond most closely to the quantiles in the source image 45 | interp_t_values = np.interp(s_quantiles, t_quantiles, t_values) 46 | 47 | return interp_t_values[bin_idx].reshape(oldshape) 48 | 49 | 50 | if __name__ == "__main__": 51 | mean_hist = np.load(MEAN_HIST_PATH) 52 | -------------------------------------------------------------------------------- /make_inpaint_images.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join 3 | import cv2 4 | import numpy as np 5 | from skimage.restoration import inpaint_biharmonic 6 | import multiprocessing as mp 7 | from multiprocessing import Pool 8 | 9 | from src.utils import make_dir 10 | from src.config import TRAIN_DIR 11 | from src.config import TEST_DIR 12 | 13 | 14 | ORIG_SIZE = (101, 101) 15 | SAVE_SIZE = (148, 148) 16 | SAVE_NAME = '148' 17 | TARGET_THRESHOLD = 0.7 18 | N_WORKERS = mp.cpu_count() 19 | 20 | 21 | make_dir(join(TRAIN_DIR, 'images'+SAVE_NAME)) 22 | make_dir(join(TRAIN_DIR, 'masks'+SAVE_NAME)) 23 | make_dir(join(TEST_DIR, 'images'+SAVE_NAME)) 24 | 25 | diff = (SAVE_SIZE - np.array(ORIG_SIZE)) 26 | pad_left = diff // 2 27 | pad_right = diff - pad_left 28 | PAD_WIDTH = ((pad_left[0], pad_right[0]), (pad_left[1], pad_right[1])) 29 | print('Pad width:', PAD_WIDTH) 30 | 31 | MASK_INPAINT = np.zeros(ORIG_SIZE, dtype=np.uint8) 32 | MASK_INPAINT = np.pad(MASK_INPAINT, PAD_WIDTH, mode='constant', constant_values=255) 33 | 34 | 35 | def inpaint_train(img_file): 36 | img = cv2.imread(join(TRAIN_DIR, 'images', img_file), cv2.IMREAD_GRAYSCALE) 37 | trg = cv2.imread(join(TRAIN_DIR, 'masks', img_file), cv2.IMREAD_GRAYSCALE) 38 | img = np.pad(img, PAD_WIDTH, mode='constant') 39 | trg = np.pad(trg, PAD_WIDTH, mode='constant') 40 | img = (inpaint_biharmonic(img, MASK_INPAINT)*255).astype(np.uint8) 41 | trg = inpaint_biharmonic(trg, MASK_INPAINT) 42 | trg = np.where(trg > TARGET_THRESHOLD, 255, 0) 43 | cv2.imwrite(join(TRAIN_DIR, 'images'+SAVE_NAME, img_file), img) 44 | cv2.imwrite(join(TRAIN_DIR, 'masks'+SAVE_NAME, img_file), trg) 45 | 46 | 47 | def inpaint_test(img_file): 48 | img = cv2.imread(join(TEST_DIR, 'images', img_file), cv2.IMREAD_GRAYSCALE) 49 | img = np.pad(img, PAD_WIDTH, mode='constant') 50 | img = (inpaint_biharmonic(img, MASK_INPAINT)*255).astype(np.uint8) 51 | cv2.imwrite(join(TEST_DIR, 'images'+SAVE_NAME, img_file), img) 52 | 53 | 54 | if __name__ == '__main__': 55 | 56 | # Train 57 | print('Start train inpaint') 58 | with Pool(processes=N_WORKERS) as pool: 59 | pool.map(inpaint_train, os.listdir(join(TRAIN_DIR, 'images'))) 60 | 61 | # Test 62 | print('Start test inpaint') 63 | with Pool(processes=N_WORKERS) as pool: 64 | pool.map(inpaint_test, os.listdir(join(TEST_DIR, 'images'))) 65 | 66 | print('Inpaint complete') 67 | -------------------------------------------------------------------------------- /pipeline/mean_submission.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | import numpy as np 3 | import pandas as pd 4 | import cv2 5 | 6 | from src import config 7 | from src.utils import RLenc, make_dir 8 | 9 | PREDICTION_DIRS = [ 10 | '/workdir/data/predictions/mos-fpn-lovasz-se-resnext50-001/' 11 | ] 12 | FOLDS = [0, 1, 2, 3, 4, 5] 13 | FOLD_DIRS = [join(p, 'fold_%d'%f) for p in PREDICTION_DIRS for f in FOLDS] 14 | 15 | PREDICTION_DIRS = [ 16 | '/workdir/data/predictions/fpn-lovasz-se-resnext50-006-after-001/', 17 | ] 18 | 19 | FOLDS = [0, 1, 2, 3, 4] 20 | FOLD_DIRS += [join(p, 'fold_%d'%f) for p in PREDICTION_DIRS for f in FOLDS] 21 | 22 | segm_thresh = 0.4 23 | prob_thresh = 0.5 24 | 25 | SAVE_NAME = 'mean-005-0.4' 26 | MEAN_PREDICTION_DIR = f'/workdir/data/predictions/{SAVE_NAME}' 27 | 28 | make_dir(join(MEAN_PREDICTION_DIR, 'masks')) 29 | 30 | 31 | def get_mean_probs_df(): 32 | probs_df_lst = [] 33 | for fold_dir in FOLD_DIRS: 34 | probs_df = pd.read_csv(join(fold_dir, 'test', 'probs.csv'), index_col='id') 35 | probs_df_lst.append(probs_df) 36 | 37 | mean_probs_df = probs_df_lst[0].copy() 38 | for probs_df in probs_df_lst[1:]: 39 | mean_probs_df.prob += probs_df.prob 40 | mean_probs_df.prob /= len(probs_df_lst) 41 | 42 | return mean_probs_df 43 | 44 | 45 | if __name__ == "__main__": 46 | print('Make average submission', FOLD_DIRS) 47 | 48 | mean_probs_df = get_mean_probs_df() 49 | sample_submition = pd.read_csv(config.SAMPLE_SUBM_PATH) 50 | 51 | for i, row in sample_submition.iterrows(): 52 | pred_name = row.id + '.png' 53 | pred_lst = [] 54 | for fold_dir in FOLD_DIRS: 55 | pred_path = join(fold_dir, 'test', pred_name) 56 | pred = cv2.imread(pred_path, cv2.IMREAD_GRAYSCALE) 57 | pred = pred / 255 58 | pred_lst.append(pred) 59 | 60 | mean_pred = np.mean(pred_lst, axis=0) 61 | prob = mean_probs_df.loc[row.id].prob 62 | 63 | pred = mean_pred > segm_thresh 64 | prob = int(prob > prob_thresh) 65 | pred = (pred * prob).astype(np.uint8) 66 | 67 | if np.all(pred == 1): 68 | pred[:] = 0 69 | print('Full mask to empty', pred_name) 70 | 71 | rle_mask = RLenc(pred) 72 | cv2.imwrite(join(MEAN_PREDICTION_DIR, 'masks', pred_name), pred * 255) 73 | row.rle_mask = rle_mask 74 | 75 | sample_submition.to_csv(join(MEAN_PREDICTION_DIR, f'{SAVE_NAME}.csv'), index=False) 76 | -------------------------------------------------------------------------------- /src/lr_scheduler.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | from argus.engine import State 4 | from argus.callbacks.callback import Callback 5 | from argus.metrics.metric import METRIC_REGISTRY 6 | 7 | 8 | class ReduceLROnPlateau(Callback): 9 | def __init__(self, 10 | monitor='val_loss', 11 | factor=0.1, 12 | patience=1, 13 | min_lr=1e-6, 14 | better='auto'): 15 | self.monitor = monitor 16 | self.factor = factor 17 | self.patience = patience 18 | self.min_lr = min_lr 19 | self.better = better 20 | 21 | if self.better == 'auto': 22 | if monitor.startswith('val_'): 23 | metric_name = self.monitor[len('val_'):] 24 | else: 25 | metric_name = self.monitor[len('train_'):] 26 | if metric_name not in METRIC_REGISTRY: 27 | raise ImportError(f"Metric '{metric_name}' not found in scope") 28 | self.better = METRIC_REGISTRY[metric_name].better 29 | assert self.better in ['min', 'max', 'auto'], \ 30 | f"Unknown better option '{self.better}'" 31 | 32 | if self.better == 'min': 33 | self.better_comp = lambda a, b: a < b 34 | self.best_value = math.inf 35 | elif self.better == 'max': 36 | self.better_comp = lambda a, b: a > b 37 | self.best_value = -math.inf 38 | 39 | self.wait = 0 40 | 41 | def start(self, state: State): 42 | self.wait = 0 43 | self.best_value = math.inf if self.better == 'min' else -math.inf 44 | 45 | def epoch_complete(self, state: State): 46 | assert self.monitor in state.metrics,\ 47 | f"Monitor '{self.monitor}' metric not found in state" 48 | current_value = state.metrics[self.monitor] 49 | if self.better_comp(current_value, self.best_value): 50 | self.best_value = current_value 51 | self.wait = 0 52 | else: 53 | self.wait += 1 54 | if self.wait >= self.patience: 55 | new_lr_lst = [] 56 | for param_group in state.model.optimizer.param_groups: 57 | param_group['lr'] *= self.factor 58 | if param_group['lr'] < self.min_lr: 59 | param_group['lr'] = self.min_lr 60 | new_lr_lst.append(param_group['lr']) 61 | 62 | state.logger.info( 63 | f"Epoch {state.epoch}: Reduce learning rate on plateau triggered, " 64 | f"'{self.monitor}' didn't improve score for {self.wait} epochs. New lr: {new_lr_lst}") 65 | self.wait = 0 66 | -------------------------------------------------------------------------------- /unused/make_mosaic_inpaint_images.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | import cv2 3 | import numpy as np 4 | from skimage.restoration import inpaint_biharmonic 5 | import multiprocessing as mp 6 | from multiprocessing import Pool 7 | 8 | from src.utils import make_dir 9 | from mosaic.mosaic_api import SaltData 10 | from src.config import DATA_DIR, TRAIN_DIR, TEST_DIR 11 | from src.transforms import CenterCrop 12 | 13 | 14 | ORIG_SIZE = (101, 101) 15 | SAVE_SIZE = (148, 148) 16 | CENTER_CROP = CenterCrop(SAVE_SIZE) 17 | SAVE_NAME = '_pazzles_6013_148' 18 | TARGET_THRESHOLD = 0.7 19 | N_WORKERS = mp.cpu_count() 20 | 21 | make_dir(join(TRAIN_DIR, 'images' + SAVE_NAME)) 22 | make_dir(join(TRAIN_DIR, 'masks' + SAVE_NAME)) 23 | make_dir(join(TEST_DIR, 'images' + SAVE_NAME)) 24 | 25 | IMAGES_DIR_NAME = "images_hist_mean" 26 | MASKS_DIR_NAME = "masks" 27 | MOSAIC_CSV_PATH = join(DATA_DIR, "mosaic", "pazzles_6013.csv") 28 | SALT_DATA = SaltData(images_dir_name=IMAGES_DIR_NAME, 29 | masks_dir_name=MASKS_DIR_NAME, 30 | mosaic_csv_path=MOSAIC_CSV_PATH) 31 | 32 | 33 | def inpaint_train(id): 34 | neighbors = SALT_DATA.get_neighbors(id) 35 | stacked_image, image_unknown_mask = SALT_DATA.get_stacked_images(neighbors, return_unknown_mask=True) 36 | center_stacked_image = CENTER_CROP(stacked_image) 37 | center_image_unknown_mask = CENTER_CROP(image_unknown_mask) 38 | if np.sum(center_image_unknown_mask): 39 | inpaint_image = inpaint_biharmonic(center_stacked_image, center_image_unknown_mask) 40 | inpaint_image = (inpaint_image * 255).astype(np.uint8) 41 | else: 42 | inpaint_image = center_stacked_image 43 | 44 | stacked_mask, mask_unknown_mask = SALT_DATA.get_stacked_masks(neighbors, return_unknown_mask=True) 45 | center_stacked_mask = CENTER_CROP(stacked_mask) 46 | center_mask_unknown_mask = CENTER_CROP(mask_unknown_mask) 47 | if np.sum(center_mask_unknown_mask): 48 | inpaint_mask = inpaint_biharmonic(center_stacked_mask, center_mask_unknown_mask) 49 | inpaint_mask= np.where(inpaint_mask > TARGET_THRESHOLD, 255, 0) 50 | else: 51 | inpaint_mask = center_stacked_mask 52 | 53 | cv2.imwrite(join(TRAIN_DIR, 'images' + SAVE_NAME, f"{id}.png"), inpaint_image) 54 | cv2.imwrite(join(TRAIN_DIR, 'masks' + SAVE_NAME, f"{id}.png"), inpaint_mask) 55 | 56 | 57 | def inpaint_test(id): 58 | neighbors = SALT_DATA.get_neighbors(id) 59 | stacked_image, image_unknown_mask = SALT_DATA.get_stacked_images(neighbors, return_unknown_mask=True) 60 | center_stacked_image = CENTER_CROP(stacked_image) 61 | center_image_unknown_mask = CENTER_CROP(image_unknown_mask) 62 | if np.sum(center_image_unknown_mask): 63 | inpaint_image = inpaint_biharmonic(center_stacked_image, center_image_unknown_mask) 64 | inpaint_image = (inpaint_image * 255).astype(np.uint8) 65 | else: 66 | inpaint_image = center_stacked_image 67 | 68 | cv2.imwrite(join(TEST_DIR, 'images' + SAVE_NAME, f"{id}.png"), inpaint_image) 69 | 70 | 71 | if __name__ == '__main__': 72 | # Train 73 | print('Start train inpaint') 74 | with Pool(processes=N_WORKERS) as pool: 75 | pool.map(inpaint_train, list(SALT_DATA.train_ids)) 76 | 77 | # Test 78 | print('Start test inpaint') 79 | with Pool(processes=N_WORKERS) as pool: 80 | pool.map(inpaint_test, list(SALT_DATA.test_ids)) 81 | 82 | print('Inpaint complete') 83 | -------------------------------------------------------------------------------- /train_folds.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | 4 | import argus 5 | from argus.callbacks import MonitorCheckpoint, EarlyStopping, LoggingToFile 6 | 7 | from torch.utils.data import DataLoader 8 | 9 | from src.dataset import SaltDataset 10 | from src.transforms import SimpleDepthTransform, SaltTransform 11 | from src.lr_scheduler import ReduceLROnPlateau 12 | from src.argus_models import SaltMetaModel 13 | from src.losses import LovaszProbLoss 14 | from src import config 15 | 16 | 17 | EXPERIMENT_NAME = 'mos-fpn-lovasz-se-resnext50-001' 18 | BATCH_SIZE = 24 19 | IMAGE_SIZE = (128, 128) 20 | OUTPUT_SIZE = (101, 101) 21 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148_mos_emb_1.csv' 22 | SAVE_DIR = f'/workdir/data/experiments/{EXPERIMENT_NAME}' 23 | FOLDS = list(range(config.N_FOLDS)) 24 | PARAMS = { 25 | 'nn_module': ('SeResnextFPNProb50', { 26 | 'num_classes': 1, 27 | 'num_channels': 3, 28 | 'final': 'logits', 29 | 'dropout_2d': 0.2, 30 | 'skip_dropout': True, 31 | 'fpn_layers': [8, 16, 32, 64, 128] 32 | }), 33 | 'loss': ('LovaszProbLoss', { 34 | 'lovasz_weight': 0.75, 35 | 'prob_weight': 0.25, 36 | }), 37 | 'prediction_transform': ('ProbOutputTransform', { 38 | 'segm_thresh': 0.0, 39 | 'prob_thresh': 0.5, 40 | }), 41 | 'optimizer': ('SGD', {'lr': 0.01, 'momentum': 0.9, 'weight_decay': 0.0001}), 42 | 'device': 'cuda' 43 | } 44 | 45 | 46 | def train_fold(save_dir, train_folds, val_folds): 47 | depth_trns = SimpleDepthTransform() 48 | train_trns = SaltTransform(IMAGE_SIZE, True, 'crop') 49 | val_trns = SaltTransform(IMAGE_SIZE, False, 'crop') 50 | train_dataset = SaltDataset(TRAIN_FOLDS_PATH, train_folds, train_trns, depth_trns) 51 | val_dataset = SaltDataset(TRAIN_FOLDS_PATH, val_folds, val_trns, depth_trns) 52 | train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, 53 | drop_last=True, num_workers=4) 54 | val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4) 55 | 56 | model = SaltMetaModel(PARAMS) 57 | 58 | callbacks = [ 59 | MonitorCheckpoint(save_dir, monitor='val_crop_iout', max_saves=3, copy_last=False), 60 | EarlyStopping(monitor='val_crop_iout', patience=100), 61 | ReduceLROnPlateau(monitor='val_crop_iout', patience=30, factor=0.64, min_lr=1e-8), 62 | LoggingToFile(os.path.join(save_dir, 'log.txt')), 63 | ] 64 | 65 | model.fit(train_loader, 66 | val_loader=val_loader, 67 | max_epochs=700, 68 | callbacks=callbacks, 69 | metrics=['crop_iout']) 70 | 71 | 72 | if __name__ == "__main__": 73 | if not os.path.exists(SAVE_DIR): 74 | os.makedirs(SAVE_DIR) 75 | else: 76 | print(f"Folder {SAVE_DIR} already exists.") 77 | 78 | with open(os.path.join(SAVE_DIR, 'source.py'), 'w') as outfile: 79 | outfile.write(open(__file__).read()) 80 | 81 | with open(os.path.join(SAVE_DIR, 'params.json'), 'w') as outfile: 82 | json.dump(PARAMS, outfile) 83 | 84 | for i in range(len(FOLDS)): 85 | val_folds = [FOLDS[i]] 86 | train_folds = FOLDS[:i] + FOLDS[i + 1:] 87 | save_fold_dir = os.path.join(SAVE_DIR, f'fold_{FOLDS[i]}') 88 | print(f"Val folds: {val_folds}, Train folds: {train_folds}") 89 | print(f"Fold save dir {save_fold_dir}") 90 | train_fold(save_fold_dir, train_folds, val_folds) 91 | -------------------------------------------------------------------------------- /unused/make_double_inpaint_images.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join 3 | import cv2 4 | import numpy as np 5 | import pandas as pd 6 | from skimage.restoration import inpaint_biharmonic 7 | import multiprocessing as mp 8 | from multiprocessing import Pool 9 | from collections import defaultdict 10 | from itertools import combinations 11 | from random import shuffle 12 | 13 | from src.utils import make_dir 14 | from src import config 15 | from src.config import TRAIN_DIR 16 | 17 | 18 | ORIG_SIZE = (101, 101) 19 | SAVE_SIZE = (109, 109) 20 | SAVE_NAME = '_double_109' 21 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148.csv' 22 | TARGET_THRESHOLD = 0.7 23 | N_WORKERS = mp.cpu_count() 24 | 25 | 26 | make_dir(join(TRAIN_DIR, 'images'+SAVE_NAME)) 27 | make_dir(join(TRAIN_DIR, 'masks'+SAVE_NAME)) 28 | 29 | diff = (SAVE_SIZE - np.array(ORIG_SIZE)) 30 | pad_left = diff // 2 31 | pad_right = diff - pad_left 32 | PAD_WIDTH = ((pad_left[0], pad_right[0]), (pad_left[0], pad_right[1])) 33 | print('Pad width:', PAD_WIDTH) 34 | 35 | MASK_INPAINT = np.zeros(ORIG_SIZE, dtype=np.uint8) 36 | MASK_INPAINT = np.pad(MASK_INPAINT, PAD_WIDTH, mode='constant', constant_values=255) 37 | MASK_INPAINT = np.concatenate([MASK_INPAINT, MASK_INPAINT], axis=1) 38 | 39 | 40 | def inpaint_double_train(img_files): 41 | img_file_1, img_file_2, fold = img_files 42 | img_1 = cv2.imread(join(TRAIN_DIR, 'images', img_file_1), cv2.IMREAD_GRAYSCALE) 43 | trg_1 = cv2.imread(join(TRAIN_DIR, 'masks', img_file_1), cv2.IMREAD_GRAYSCALE) 44 | img_2 = cv2.imread(join(TRAIN_DIR, 'images', img_file_2), cv2.IMREAD_GRAYSCALE) 45 | trg_2 = cv2.imread(join(TRAIN_DIR, 'masks', img_file_2), cv2.IMREAD_GRAYSCALE) 46 | 47 | if not np.sum(img_1) or not np.sum(img_2): 48 | return 49 | 50 | img_1 = np.pad(img_1, PAD_WIDTH, mode='constant') 51 | trg_1 = np.pad(trg_1, PAD_WIDTH, mode='constant') 52 | img_2 = np.pad(img_2, PAD_WIDTH, mode='constant') 53 | trg_2 = np.pad(trg_2, PAD_WIDTH, mode='constant') 54 | img = np.concatenate([img_1, img_2], axis=1) 55 | trg = np.concatenate([trg_1, trg_2], axis=1) 56 | img = (inpaint_biharmonic(img, MASK_INPAINT)*255).astype(np.uint8) 57 | trg = inpaint_biharmonic(trg, MASK_INPAINT) 58 | trg = np.where(trg > TARGET_THRESHOLD, 255, 0) 59 | 60 | img_file = img_file_1[:-4] + '_' + img_file_2[:-4] + '.png' 61 | make_dir(join(TRAIN_DIR, 'images' + SAVE_NAME, 'fold_%d'%fold)) 62 | make_dir(join(TRAIN_DIR, 'masks'+SAVE_NAME, 'fold_%d'%fold)) 63 | cv2.imwrite(join(TRAIN_DIR, 'images'+SAVE_NAME, 'fold_%d'%fold, img_file), img) 64 | cv2.imwrite(join(TRAIN_DIR, 'masks'+SAVE_NAME, 'fold_%d'%fold, img_file), trg) 65 | return img, trg, img_file 66 | 67 | 68 | if __name__ == '__main__': 69 | 70 | train_folds_df = pd.read_csv(TRAIN_FOLDS_PATH) 71 | fold2names = defaultdict(list) 72 | for i, row in train_folds_df.iterrows(): 73 | name = os.path.basename(row.image_path) 74 | fold2names[row.fold].append(name) 75 | 76 | double_image_lst = [] 77 | 78 | for fold in range(config.N_FOLDS): 79 | print(fold) 80 | fold_names = list(combinations(fold2names[fold], 2)) 81 | shuffle(fold_names) 82 | for name1, name2 in fold_names[:20000]: 83 | double_image_lst.append((name1, name2, fold)) 84 | 85 | shuffle(double_image_lst) 86 | 87 | with Pool(processes=N_WORKERS) as pool: 88 | pool.map(inpaint_double_train, double_image_lst) 89 | 90 | print('Inpaint complete') 91 | -------------------------------------------------------------------------------- /pipeline/mos-fpn-lovasz-se-resnext50-001/train_folds.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | 4 | import argus 5 | from argus.callbacks import MonitorCheckpoint, EarlyStopping, LoggingToFile 6 | 7 | from torch.utils.data import DataLoader 8 | 9 | from src.dataset import SaltDataset 10 | from src.transforms import SimpleDepthTransform, SaltTransform 11 | from src.lr_scheduler import ReduceLROnPlateau 12 | from src.argus_models import SaltMetaModel 13 | from src.losses import LovaszProbLoss 14 | from src import config 15 | 16 | 17 | EXPERIMENT_NAME = 'mos-fpn-lovasz-se-resnext50-001' 18 | BATCH_SIZE = 24 19 | IMAGE_SIZE = (128, 128) 20 | OUTPUT_SIZE = (101, 101) 21 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148_mos_emb_1.csv' 22 | SAVE_DIR = f'/workdir/data/experiments/{EXPERIMENT_NAME}' 23 | FOLDS = list(range(config.N_FOLDS)) 24 | PARAMS = { 25 | 'nn_module': ('SeResnextFPNProb50', { 26 | 'num_classes': 1, 27 | 'num_channels': 3, 28 | 'final': 'logits', 29 | 'dropout_2d': 0.2, 30 | 'skip_dropout': True, 31 | 'fpn_layers': [8, 16, 32, 64, 128] 32 | }), 33 | 'loss': ('LovaszProbLoss', { 34 | 'lovasz_weight': 0.75, 35 | 'prob_weight': 0.25, 36 | }), 37 | 'prediction_transform': ('ProbOutputTransform', { 38 | 'segm_thresh': 0.0, 39 | 'prob_thresh': 0.5, 40 | }), 41 | 'optimizer': ('SGD', {'lr': 0.01, 'momentum': 0.9, 'weight_decay': 0.0001}), 42 | 'device': 'cuda' 43 | } 44 | 45 | 46 | def train_fold(save_dir, train_folds, val_folds): 47 | depth_trns = SimpleDepthTransform() 48 | train_trns = SaltTransform(IMAGE_SIZE, True, 'crop') 49 | val_trns = SaltTransform(IMAGE_SIZE, False, 'crop') 50 | train_dataset = SaltDataset(TRAIN_FOLDS_PATH, train_folds, train_trns, depth_trns) 51 | val_dataset = SaltDataset(TRAIN_FOLDS_PATH, val_folds, val_trns, depth_trns) 52 | train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, 53 | drop_last=True, num_workers=4) 54 | val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4) 55 | 56 | model = SaltMetaModel(PARAMS) 57 | 58 | callbacks = [ 59 | MonitorCheckpoint(save_dir, monitor='val_crop_iout', max_saves=3, copy_last=False), 60 | EarlyStopping(monitor='val_crop_iout', patience=100), 61 | ReduceLROnPlateau(monitor='val_crop_iout', patience=30, factor=0.64, min_lr=1e-8), 62 | LoggingToFile(os.path.join(save_dir, 'log.txt')), 63 | ] 64 | 65 | model.fit(train_loader, 66 | val_loader=val_loader, 67 | max_epochs=700, 68 | callbacks=callbacks, 69 | metrics=['crop_iout']) 70 | 71 | 72 | if __name__ == "__main__": 73 | if not os.path.exists(SAVE_DIR): 74 | os.makedirs(SAVE_DIR) 75 | else: 76 | print(f"Folder {SAVE_DIR} already exists.") 77 | 78 | with open(os.path.join(SAVE_DIR, 'source.py'), 'w') as outfile: 79 | outfile.write(open(__file__).read()) 80 | 81 | with open(os.path.join(SAVE_DIR, 'params.json'), 'w') as outfile: 82 | json.dump(PARAMS, outfile) 83 | 84 | for i in range(len(FOLDS)): 85 | val_folds = [FOLDS[i]] 86 | train_folds = FOLDS[:i] + FOLDS[i + 1:] 87 | save_fold_dir = os.path.join(SAVE_DIR, f'fold_{FOLDS[i]}') 88 | print(f"Val folds: {val_folds}, Train folds: {train_folds}") 89 | print(f"Fold save dir {save_fold_dir}") 90 | train_fold(save_fold_dir, train_folds, val_folds) 91 | -------------------------------------------------------------------------------- /pipeline/fpn-lovasz-se-resnext50-006/train_folds.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | 4 | import argus 5 | from argus.callbacks import MonitorCheckpoint, EarlyStopping, LoggingToFile 6 | 7 | from torch.utils.data import DataLoader 8 | 9 | from src.dataset import SaltDataset 10 | from src.transforms import SimpleDepthTransform, SaltTransform 11 | from src.lr_scheduler import ReduceLROnPlateau 12 | from src.argus_models import SaltMetaModel 13 | from src.losses import LovaszProbLoss 14 | from src import config 15 | 16 | from src.nick_zoo.resnet_blocks import resnet34 17 | 18 | 19 | EXPERIMENT_NAME = 'fpn-lovasz-se-resnext50-006' 20 | BATCH_SIZE = 24 21 | IMAGE_SIZE = (128, 128) 22 | OUTPUT_SIZE = (101, 101) 23 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148.csv' 24 | SAVE_DIR = f'/workdir/data/experiments/{EXPERIMENT_NAME}' 25 | N_FOLDS = 5 26 | FOLDS = list(range(N_FOLDS)) 27 | PARAMS = { 28 | 'nn_module': ('SeResnextFPNProb50', { 29 | 'num_classes': 1, 30 | 'num_channels': 3, 31 | 'final': 'logits', 32 | 'dropout_2d': 0.2, 33 | 'skip_dropout': True, 34 | 'fpn_layers': [8, 16, 32, 64, 128] 35 | }), 36 | 'loss': ('LovaszProbLoss', { 37 | 'lovasz_weight': 0.75, 38 | 'prob_weight': 0.25, 39 | }), 40 | 'prediction_transform': ('ProbOutputTransform', { 41 | 'segm_thresh': 0.0, 42 | 'prob_thresh': 0.5, 43 | }), 44 | 'optimizer': ('SGD', {'lr': 0.01, 'momentum': 0.9, 'weight_decay': 0.0001}), 45 | 'device': 'cuda' 46 | } 47 | 48 | 49 | def train_fold(save_dir, train_folds, val_folds): 50 | depth_trns = SimpleDepthTransform() 51 | train_trns = SaltTransform(IMAGE_SIZE, True, 'crop') 52 | val_trns = SaltTransform(IMAGE_SIZE, False, 'crop') 53 | train_dataset = SaltDataset(TRAIN_FOLDS_PATH, train_folds, train_trns, depth_trns) 54 | val_dataset = SaltDataset(TRAIN_FOLDS_PATH, val_folds, val_trns, depth_trns) 55 | train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, 56 | drop_last=True, num_workers=4) 57 | val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4) 58 | 59 | model = SaltMetaModel(PARAMS) 60 | 61 | callbacks = [ 62 | MonitorCheckpoint(save_dir, monitor='val_crop_iout', max_saves=3, copy_last=False), 63 | EarlyStopping(monitor='val_crop_iout', patience=100), 64 | ReduceLROnPlateau(monitor='val_crop_iout', patience=30, factor=0.64, min_lr=1e-8), 65 | LoggingToFile(os.path.join(save_dir, 'log.txt')), 66 | ] 67 | 68 | model.fit(train_loader, 69 | val_loader=val_loader, 70 | max_epochs=700, 71 | callbacks=callbacks, 72 | metrics=['crop_iout']) 73 | 74 | 75 | if __name__ == "__main__": 76 | if not os.path.exists(SAVE_DIR): 77 | os.makedirs(SAVE_DIR) 78 | else: 79 | print(f"Folder {SAVE_DIR} already exists.") 80 | 81 | with open(os.path.join(SAVE_DIR, 'source.py'), 'w') as outfile: 82 | outfile.write(open(__file__).read()) 83 | 84 | with open(os.path.join(SAVE_DIR, 'params.json'), 'w') as outfile: 85 | json.dump(PARAMS, outfile) 86 | 87 | for i in range(len(FOLDS)): 88 | val_folds = [FOLDS[i]] 89 | train_folds = FOLDS[:i] + FOLDS[i + 1:] 90 | save_fold_dir = os.path.join(SAVE_DIR, f'fold_{FOLDS[i]}') 91 | print(f"Val folds: {val_folds}, Train folds: {train_folds}") 92 | print(f"Fold save dir {save_fold_dir}") 93 | train_fold(save_fold_dir, train_folds, val_folds) 94 | -------------------------------------------------------------------------------- /src/dataset.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import os 3 | from os.path import join 4 | import numpy as np 5 | import pandas as pd 6 | 7 | from torch.utils.data import Dataset 8 | 9 | from src.transforms import SimpleDepthTransform 10 | from src.config import TEST_DIR 11 | 12 | 13 | def get_samples(train_folds_path, folds): 14 | images_lst = [] 15 | target_lst = [] 16 | depth_lst = [] 17 | 18 | train_folds_df = pd.read_csv(train_folds_path) 19 | 20 | for i, row in train_folds_df.iterrows(): 21 | if row.fold not in folds: 22 | continue 23 | 24 | image = cv2.imread(row.image_path, cv2.IMREAD_GRAYSCALE) 25 | if image is None: 26 | raise FileNotFoundError(f"Image not found {row.image_path}") 27 | mask = cv2.imread(row.mask_path, cv2.IMREAD_GRAYSCALE) 28 | if mask is None: 29 | raise FileNotFoundError(f"Mask not found {row.mask_path}") 30 | images_lst.append(image) 31 | target_lst.append(mask) 32 | depth_lst.append(row.z) 33 | 34 | return images_lst, target_lst, depth_lst 35 | 36 | 37 | class SaltDataset(Dataset): 38 | def __init__(self, train_folds_path, folds, 39 | transform=None, 40 | depth_transform=None): 41 | super().__init__() 42 | self.train_folds_path = train_folds_path 43 | self.folds = folds 44 | self.transform = transform 45 | if depth_transform is None: 46 | self.depth_transform = SimpleDepthTransform() 47 | else: 48 | self.depth_transform = depth_transform 49 | 50 | self.images_lst, self.target_lst, self.depth_lst = \ 51 | get_samples(train_folds_path, folds) 52 | 53 | def __len__(self): 54 | return len(self.images_lst) 55 | 56 | def __getitem__(self, idx): 57 | image = self.images_lst[idx] 58 | depth = self.depth_lst[idx] 59 | target = self.target_lst[idx] 60 | 61 | input = self.depth_transform(image, depth) 62 | 63 | if self.transform is not None: 64 | input, target = self.transform(input, target) 65 | 66 | return input, target 67 | 68 | 69 | def get_test_samples(test_images_dir): 70 | images_lst = [] 71 | depth_lst = [] 72 | 73 | for image_name in os.listdir(test_images_dir): 74 | image_path = join(test_images_dir, image_name) 75 | image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) 76 | if np.sum(image): # skip black images 77 | images_lst.append(image) 78 | depth_lst.append(0) # TODO: load depth 79 | 80 | return images_lst, depth_lst 81 | 82 | 83 | class SaltTestDataset(Dataset): 84 | def __init__(self, test_dir, 85 | transform=None, 86 | depth_transform=None): 87 | super().__init__() 88 | self.test_dir = test_dir 89 | self.transform = transform 90 | if depth_transform is None: 91 | self.depth_transform = SimpleDepthTransform() 92 | else: 93 | self.depth_transform = depth_transform 94 | 95 | self.images_lst, self.depth_lst = \ 96 | get_test_samples(test_dir) 97 | 98 | def __len__(self): 99 | return len(self.images_lst) 100 | 101 | def __getitem__(self, idx): 102 | image = self.images_lst[idx] 103 | depth = self.depth_lst[idx] 104 | 105 | input = self.depth_transform(image, depth) 106 | 107 | if self.transform is not None: 108 | input = self.transform(input) 109 | 110 | return input 111 | -------------------------------------------------------------------------------- /src/metrics.py: -------------------------------------------------------------------------------- 1 | # From https://github.com/neptune-ml/open-solution-salt-detection/blob/master/src/metrics.py 2 | 3 | import numpy as np 4 | from pycocotools import mask as cocomask 5 | 6 | from argus.metrics import Metric 7 | 8 | from src.transforms import CenterCrop 9 | 10 | 11 | def rle_from_binary(prediction): 12 | prediction = np.asfortranarray(prediction) 13 | return cocomask.encode(prediction) 14 | 15 | 16 | def binary_from_rle(rle): 17 | return cocomask.decode(rle) 18 | 19 | 20 | def get_segmentations(labeled): 21 | nr_true = labeled.max() 22 | segmentations = [] 23 | for i in range(1, nr_true + 1): 24 | msk = labeled == i 25 | segmentation = rle_from_binary(msk.astype('uint8')) 26 | segmentation['counts'] = segmentation['counts'].decode("UTF-8") 27 | segmentations.append(segmentation) 28 | return segmentations 29 | 30 | 31 | def iou(gt, pred): 32 | gt[gt > 0] = 1. 33 | pred[pred > 0] = 1. 34 | intersection = gt * pred 35 | union = gt + pred 36 | union[union > 0] = 1. 37 | intersection = np.sum(intersection) 38 | union = np.sum(union) 39 | if union == 0: 40 | union = 1e-09 41 | return intersection / union 42 | 43 | 44 | def compute_ious(gt, predictions): 45 | gt_ = get_segmentations(gt) 46 | predictions_ = get_segmentations(predictions) 47 | 48 | if len(gt_) == 0 and len(predictions_) == 0: 49 | return np.ones((1, 1)) 50 | elif len(gt_) != 0 and len(predictions_) == 0: 51 | return np.zeros((1, 1)) 52 | else: 53 | iscrowd = [0 for _ in predictions_] 54 | ious = cocomask.iou(gt_, predictions_, iscrowd) 55 | if not np.array(ious).size: 56 | ious = np.zeros((1, 1)) 57 | return ious 58 | 59 | 60 | def compute_precision_at(ious, threshold): 61 | mx1 = np.max(ious, axis=0) 62 | mx2 = np.max(ious, axis=1) 63 | tp = np.sum(mx2 >= threshold) 64 | fp = np.sum(mx2 < threshold) 65 | fn = np.sum(mx1 < threshold) 66 | return float(tp) / (tp + fp + fn) 67 | 68 | 69 | def compute_eval_metric(gt, predictions): 70 | thresholds = [0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95] 71 | ious = compute_ious(gt, predictions) 72 | precisions = [compute_precision_at(ious, th) for th in thresholds] 73 | return sum(precisions) / len(precisions) 74 | 75 | 76 | class SaltIOUT(Metric): 77 | name = 'iout' 78 | better = 'max' 79 | 80 | def reset(self): 81 | self.iout_sum = 0 82 | self.count = 0 83 | 84 | def update(self, step_output: dict): 85 | preds = step_output['prediction'].cpu().numpy().astype(np.uint8) 86 | trgs = step_output['target'].cpu().numpy().astype(np.uint8) 87 | 88 | for i in range(trgs.shape[0]): 89 | pred = preds[i][0] 90 | trg = trgs[i][0] 91 | self.iout_sum += compute_eval_metric(trg, pred) 92 | self.count += 1 93 | 94 | def compute(self): 95 | return self.iout_sum / self.count 96 | 97 | 98 | class SaltCropIOUT(Metric): 99 | name = 'crop_iout' 100 | better = 'max' 101 | 102 | def __init__(self, size=(101, 101)): 103 | super().__init__() 104 | self.crop_trns = CenterCrop(size) 105 | 106 | def reset(self): 107 | self.iout_sum = 0 108 | self.count = 0 109 | 110 | def update(self, step_output: dict): 111 | preds = step_output['prediction'].cpu().numpy().astype(np.uint8) 112 | trgs = step_output['target'].cpu().numpy().astype(np.uint8) 113 | 114 | for i in range(trgs.shape[0]): 115 | pred = preds[i][0] 116 | trg = trgs[i][0] 117 | pred, trg = self.crop_trns(pred, trg) 118 | self.iout_sum += compute_eval_metric(trg, pred) 119 | self.count += 1 120 | 121 | def compute(self): 122 | return self.iout_sum / self.count 123 | -------------------------------------------------------------------------------- /random_search.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import random 4 | import numpy as np 5 | from pprint import pprint 6 | 7 | from argus.callbacks import MonitorCheckpoint, EarlyStopping, LoggingToFile 8 | 9 | from torch.utils.data import DataLoader 10 | 11 | from src.dataset import SaltDataset 12 | from src.transforms import SimpleDepthTransform, SaltTransform 13 | from src.lr_scheduler import ReduceLROnPlateau 14 | from src.argus_models import SaltMetaModel 15 | 16 | from src.nick_zoo.resnet_blocks import resnet34 17 | 18 | 19 | SAVE_DIR = 'random-search-flex-fpn-resnet34-001' 20 | VAL_FOLDS = [0] 21 | TRAIN_FOLDS = [1, 2, 3, 4] 22 | START_FROM = 0 23 | BATCH_SIZE = 32 24 | IMAGE_SIZE = (128, 128) 25 | OUTPUT_SIZE = (101, 101) 26 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148.csv' 27 | 28 | 29 | if __name__ == "__main__": 30 | for i in range(START_FROM, 1000): 31 | experiment_dir = f'/workdir/data/experiments/{SAVE_DIR}/{i:03}' 32 | np.random.seed(i) 33 | random.seed(i) 34 | random_params = { 35 | 'dropout': float(np.random.uniform(0.0, 1.0)), 36 | 'fb_weight': float(np.random.uniform(0.5, 1.0)), 37 | 'fb_beta': float(np.random.choice([0.5, 1, 2])), 38 | 'bce_weight': float(np.random.uniform(0.5, 1.0)), 39 | 'prob_weight': float(np.random.uniform(0.5, 1.0)) 40 | } 41 | 42 | params = { 43 | 'nn_module': ('UNetFPNFlexProb', { 44 | 'num_classes': 1, 45 | 'num_channels': 3, 46 | 'blocks': resnet34, 47 | 'final': 'sigmoid', 48 | 'dropout_2d': random_params['dropout'], 49 | 'is_deconv': True, 50 | 'deflation': 4, 51 | 'use_first_pool': False, 52 | 'skip_dropout': True, 53 | 'pretrain': 'resnet34', 54 | 'pretrain_layers': [True for _ in range(5)], 55 | 'fpn_layers': [16, 32, 64, 128] 56 | }), 57 | 'loss': ('FbBceProbLoss', { 58 | 'fb_weight': random_params['fb_weight'], 59 | 'fb_beta': random_params['fb_beta'], 60 | 'bce_weight': random_params['bce_weight'], 61 | 'prob_weight': random_params['prob_weight'] 62 | }), 63 | 'prediction_transform': ('ProbOutputTransform', { 64 | 'segm_thresh': 0.5, 65 | 'prob_thresh': 0.5, 66 | }), 67 | 'optimizer': ('Adam', {'lr': 0.0001}), 68 | 'device': 'cuda' 69 | } 70 | pprint(params) 71 | 72 | depth_trns = SimpleDepthTransform() 73 | train_trns = SaltTransform(IMAGE_SIZE, True, 'crop') 74 | val_trns = SaltTransform(IMAGE_SIZE, False, 'crop') 75 | train_dataset = SaltDataset(TRAIN_FOLDS_PATH, TRAIN_FOLDS, train_trns, depth_trns) 76 | val_dataset = SaltDataset(TRAIN_FOLDS_PATH, VAL_FOLDS, val_trns, depth_trns) 77 | train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, 78 | drop_last=True, num_workers=4) 79 | val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4) 80 | 81 | model = SaltMetaModel(params) 82 | 83 | callbacks = [ 84 | MonitorCheckpoint(experiment_dir, monitor='val_crop_iout', max_saves=1, copy_last=False), 85 | EarlyStopping(monitor='val_crop_iout', patience=100), 86 | ReduceLROnPlateau(monitor='val_crop_iout', patience=30, factor=0.7, min_lr=1e-8), 87 | LoggingToFile(os.path.join(experiment_dir, 'log.txt')) 88 | ] 89 | 90 | with open(os.path.join(experiment_dir, 'random_params.json'), 'w') as outfile: 91 | json.dump(random_params, outfile) 92 | 93 | model.fit(train_loader, 94 | val_loader=val_loader, 95 | max_epochs=600, 96 | callbacks=callbacks, 97 | metrics=['crop_iout']) 98 | -------------------------------------------------------------------------------- /src/losses.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | from src.lovasz import lovasz_hinge 5 | 6 | 7 | # Source https://github.com/NikolasEnt/Lyft-Perception-Challenge/blob/master/loss.py 8 | def fb_loss(preds, trues, beta): 9 | smooth = 1e-4 10 | beta2 = beta*beta 11 | batch = preds.size(0) 12 | classes = preds.size(1) 13 | preds = preds.view(batch, classes, -1) 14 | trues = trues.view(batch, classes, -1) 15 | weights = torch.clamp(trues.sum(-1), 0., 1.) 16 | TP = (preds * trues).sum(2) 17 | FP = (preds * (1-trues)).sum(2) 18 | FN = ((1-preds) * trues).sum(2) 19 | Fb = ((1+beta2) * TP + smooth)/((1+beta2) * TP + beta2 * FN + FP + smooth) 20 | Fb = Fb * weights 21 | score = Fb.sum() / (weights.sum() + smooth) 22 | return torch.clamp(score, 0., 1.) 23 | 24 | 25 | class FBLoss(nn.Module): 26 | def __init__(self, beta=1): 27 | super().__init__() 28 | self.beta = beta 29 | 30 | def forward(self, output, target): 31 | return 1 - fb_loss(output, target, self.beta) 32 | 33 | 34 | class FbBceProbLoss(nn.Module): 35 | def __init__(self, fb_weight=0.33, fb_beta=1, bce_weight=0.33, prob_weight=0.33, min_mask_pixels=0): 36 | super().__init__() 37 | self.fb_weight = fb_weight 38 | self.bce_weight = bce_weight 39 | self.prob_weight = prob_weight 40 | self.min_mask_pixels = min_mask_pixels 41 | 42 | self.fb_loss = FBLoss(beta=fb_beta) 43 | self.bce_loss = nn.BCELoss() 44 | self.prob_loss = nn.BCELoss() 45 | 46 | def forward(self, output, target): 47 | segm, prob_pred = output 48 | if self.fb_weight > 0: 49 | fb = self.fb_loss(segm, target) * self.fb_weight 50 | else: 51 | fb = 0 52 | 53 | if self.bce_weight > 0: 54 | bce = self.bce_loss(segm, target) * self.bce_weight 55 | else: 56 | bce = 0 57 | 58 | prob_trg = target.view(target.size(0), -1).sum(dim=1) > self.min_mask_pixels 59 | prob_trg = prob_trg.to(torch.float32) 60 | if self.prob_weight > 0: 61 | prob = self.prob_loss(prob_pred, prob_trg) * self.prob_weight 62 | else: 63 | prob = 0 64 | 65 | return fb + bce + prob 66 | 67 | 68 | class ConsistencyLoss(nn.Module): 69 | def __init__(self, segm_weight=0.33, prob_weight=0.33): 70 | super().__init__() 71 | self.segm_weight = segm_weight 72 | self.prob_weight = prob_weight 73 | 74 | self.segm_loss = nn.BCELoss() 75 | self.prob_loss = nn.BCELoss() 76 | 77 | def forward(self, student_pred, teacher_pred): 78 | student_segm, student_prob = student_pred 79 | teacher_segm, teacher_prob = teacher_pred 80 | 81 | if self.segm_weight > 0: 82 | segm = self.segm_loss(student_segm, teacher_segm) * self.segm_weight 83 | else: 84 | segm = 0 85 | 86 | if self.prob_weight > 0: 87 | prob = self.prob_loss(student_prob, teacher_prob) * self.prob_weight 88 | else: 89 | prob = 0 90 | 91 | return segm + prob 92 | 93 | 94 | class LovaszProbLoss(nn.Module): 95 | def __init__(self, lovasz_weight=0.5, prob_weight=0.5, min_mask_pixels=0): 96 | super().__init__() 97 | self.lovasz_weight = lovasz_weight 98 | self.prob_weight = prob_weight 99 | self.min_mask_pixels = min_mask_pixels 100 | self.prob_loss = nn.BCELoss() 101 | 102 | def forward(self, output, target): 103 | segm, prob_pred = output 104 | 105 | prob_trg = target.view(target.size(0), -1).sum(dim=1) > self.min_mask_pixels 106 | prob_trg = prob_trg.to(torch.float32) 107 | if self.prob_weight > 0: 108 | prob = self.prob_loss(prob_pred, prob_trg) * self.prob_weight 109 | else: 110 | prob = 0 111 | 112 | if self.lovasz_weight > 0: 113 | lovasz = lovasz_hinge(segm.squeeze(1), target.squeeze(1)) \ 114 | * self.lovasz_weight 115 | else: 116 | lovasz = 0 117 | 118 | return prob + lovasz 119 | -------------------------------------------------------------------------------- /after_train_folds.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join 3 | import re 4 | import math 5 | import torch 6 | 7 | import argus 8 | from argus.callbacks import MonitorCheckpoint, LoggingToFile 9 | from argus import load_model 10 | 11 | from torch.utils.data import DataLoader 12 | 13 | from src.dataset import SaltDataset 14 | from src.transforms import SimpleDepthTransform, SaltTransform 15 | from src.argus_models import SaltMetaModel 16 | from src.losses import LovaszProbLoss 17 | from src import config 18 | 19 | 20 | def get_best_model_path(dir_path): 21 | model_scores = [] 22 | for model_name in os.listdir(dir_path): 23 | score = re.search(r'-(\d+(?:\.\d+)?).pth', model_name) 24 | if score is not None: 25 | score = score.group(0)[1:-4] 26 | model_scores.append((model_name, score)) 27 | model_score = sorted(model_scores, key=lambda x: x[1]) 28 | best_model_name = model_score[-1][0] 29 | best_model_path = os.path.join(dir_path, best_model_name) 30 | return best_model_path 31 | 32 | 33 | BASE_EXPERIMENT_NAME = 'mos-fpn-lovasz-se-resnext50-001' 34 | EXPERIMENT_NAME = 'mos-fpn-lovasz-se-resnext50-001-after-001' 35 | FOLDS = list(range(config.N_FOLDS)) 36 | BATCH_SIZE = 16 37 | IMAGE_SIZE = (128, 128) 38 | OUTPUT_SIZE = (101, 101) 39 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148_mos_emb_1.csv' 40 | LR = 0.005 41 | SAVE_DIR = f'/workdir/data/experiments/{EXPERIMENT_NAME}' 42 | 43 | 44 | class CosineAnnealingLR: 45 | def __init__(self, base_lr, T_max, eta_min=0.): 46 | self.T_max = T_max 47 | self.eta_min = eta_min 48 | self.base_lr = base_lr 49 | 50 | def __call__(self, epoch): 51 | return self.eta_min + (self.base_lr - self.eta_min) \ 52 | * (1 + math.cos(math.pi * (epoch % self.T_max) / self.T_max)) / 2 53 | 54 | 55 | cos_ann = CosineAnnealingLR(LR, 50, eta_min=LR*0.1) 56 | 57 | 58 | @argus.callbacks.on_epoch_start 59 | def update_lr(state: argus.engine.State): 60 | lr = cos_ann(state.epoch) 61 | state.model.set_lr(lr) 62 | state.logger.info(f"Set lr: {lr}") 63 | 64 | 65 | def train_fold(save_dir, train_folds, val_folds, model_path): 66 | depth_trns = SimpleDepthTransform() 67 | train_trns = SaltTransform(IMAGE_SIZE, True, 'crop') 68 | val_trns = SaltTransform(IMAGE_SIZE, False, 'crop') 69 | train_dataset = SaltDataset(TRAIN_FOLDS_PATH, train_folds, train_trns, depth_trns) 70 | val_dataset = SaltDataset(TRAIN_FOLDS_PATH, val_folds, val_trns, depth_trns) 71 | train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, 72 | drop_last=True, num_workers=8) 73 | val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=8) 74 | 75 | model = load_model(model_path) 76 | model.loss.lovasz_weight = 0.5 77 | model.loss.prob_weight = 0.5 78 | 79 | callbacks = [ 80 | MonitorCheckpoint(save_dir, monitor='val_crop_iout', max_saves=3, copy_last=False), 81 | LoggingToFile(os.path.join(save_dir, 'log.txt')), 82 | update_lr 83 | ] 84 | 85 | model.fit(train_loader, 86 | val_loader=val_loader, 87 | max_epochs=500, 88 | callbacks=callbacks, 89 | metrics=['crop_iout']) 90 | 91 | 92 | if __name__ == "__main__": 93 | if not os.path.exists(SAVE_DIR): 94 | os.makedirs(SAVE_DIR) 95 | else: 96 | print(f"Folder {SAVE_DIR} already exists.") 97 | 98 | with open(os.path.join(SAVE_DIR, 'source.py'), 'w') as outfile: 99 | outfile.write(open(__file__).read()) 100 | 101 | for i in range(len(FOLDS)): 102 | val_folds = [FOLDS[i]] 103 | train_folds = FOLDS[:i] + FOLDS[i + 1:] 104 | save_fold_dir = os.path.join(SAVE_DIR, f'fold_{FOLDS[i]}') 105 | print(f"Val folds: {val_folds}, Train folds: {train_folds}") 106 | print(f"Fold save dir {save_fold_dir}") 107 | model_path = get_best_model_path(join('/workdir/data/experiments', BASE_EXPERIMENT_NAME, 'fold_%d' % i)) 108 | print(f'Base model path: {model_path}') 109 | train_fold(save_fold_dir, train_folds, val_folds, model_path) 110 | -------------------------------------------------------------------------------- /src/nick_zoo/unet_flex_fpn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from src.nick_zoo.unet_flex import UNetFlexProb 5 | from src.nick_zoo.resnet_blocks import resnet34 6 | from src.nick_zoo.layers import FPNBlock 7 | 8 | nonlinearity = nn.ReLU 9 | 10 | 11 | class UNetFPNFlexProb(UNetFlexProb): 12 | def __init__(self, num_classes, num_channels=3, blocks=resnet34, 13 | final='softmax', dropout_2d=0.5, deflation=4, 14 | is_deconv=True, first_last_conv=(64, 32), 15 | skip_dropout=False, xavier=False, pretrain=None, 16 | pretrain_layers=None, use_first_pool=True, fpn_layers=None): 17 | super(UNetFPNFlexProb, self).__init__(num_classes, num_channels, 18 | blocks, final, dropout_2d, deflation, 19 | is_deconv, first_last_conv, skip_dropout, 20 | xavier, pretrain, pretrain_layers, 21 | use_first_pool) 22 | 23 | if fpn_layers is None: 24 | fpn_layers = [128 for _ in range(len(self.encoders_n))] 25 | else: 26 | assert len(self.encoders_n) == len(fpn_layers), \ 27 | "Incorrect number of FPN blocks" 28 | 29 | n_blocks = len(fpn_layers) 30 | c_scale = 1 31 | 32 | self.fpn_blocks = nn.ModuleList([]) 33 | for i in range(n_blocks): 34 | self.fpn_blocks.append(FPNBlock(self.encoders_n[i], 35 | fpn_layers[i], c_scale)) 36 | c_scale = c_scale * 2 37 | 38 | fpn_filters = sum(fpn_layers) 39 | # Final classifier 40 | if self.use_first_pool: 41 | if is_deconv: 42 | self.finalupscale = nn.ConvTranspose2d(fpn_filters, 43 | first_last_conv[1], 3, 44 | stride=2) 45 | self.finalconv2 = nn.Conv2d(first_last_conv[1], first_last_conv[1], 3) 46 | else: 47 | self.finalupscale = nn.Upsample(scale_factor=2, mode='bilinear', 48 | align_corners=False) 49 | self.finalconv2 = nn.Conv2d(fpn_filters, first_last_conv[1], 3) 50 | else: 51 | self.finalconv2 = nn.Conv2d(fpn_filters, first_last_conv[1], 3, padding=1) 52 | 53 | self.prob_bn = nn.BatchNorm2d(self.encoders_n[-1]) 54 | 55 | def forward(self, x): 56 | # Encoder 57 | x = self.firstconv(x) 58 | x = self.firstbn(x) 59 | x = self.firstrelu(x) 60 | if self.use_first_pool: 61 | x = self.firstmaxpool(x) 62 | 63 | enc = [] 64 | for encoder in self.encoders: 65 | x = encoder(x) 66 | enc.append(x) 67 | 68 | # Decoder with Skip Connections 69 | dec = [] 70 | for i in range(len(self.decoders) - 1, 0, -1): 71 | x = self.decoders[i](x) 72 | if self.skip_dropouts is not None: 73 | x = torch.cat([x, self.skip_dropouts[i - 1](enc[i - 1])], dim=1) 74 | else: 75 | x = torch.cat([x, enc[i - 1]], dim=1) 76 | dec.append(self.fpn_blocks[i](x)) 77 | 78 | x = self.decoders[0](x) 79 | dec.append(self.fpn_blocks[0](x)) 80 | # Concat the FPN 81 | x = torch.cat(dec, dim=1) 82 | 83 | # Final Classification 84 | if self.use_first_pool: 85 | x = self.finalupscale(x) 86 | x = self.finalrelu1(x) 87 | if self.finaldropout is not None: 88 | x = self.finaldropout(x) 89 | x = self.finalconv2(x) 90 | x = self.finalrelu2(x) 91 | x = self.finalconv3(x) 92 | if self.final is not None: 93 | x = self.final(x) 94 | 95 | prob = self.prob_bn(enc[-1]) 96 | prob = self.prob_avg_pool(prob) 97 | prob = prob.view(prob.size(0), -1) 98 | prob = self.prob_fc(prob) 99 | prob = self.sigmoid(prob) 100 | prob = prob.view(-1) 101 | return x, prob 102 | -------------------------------------------------------------------------------- /pipeline/fpn-lovasz-se-resnext50-006-after-001/train_folds.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join 3 | import re 4 | import math 5 | import torch 6 | 7 | import argus 8 | from argus.callbacks import MonitorCheckpoint, LoggingToFile 9 | from argus import load_model 10 | 11 | from torch.utils.data import DataLoader 12 | 13 | from src.dataset import SaltDataset 14 | from src.transforms import SimpleDepthTransform, SaltTransform 15 | from src.argus_models import SaltMetaModel 16 | from src.losses import LovaszProbLoss 17 | from src import config 18 | 19 | 20 | def get_best_model_path(dir_path): 21 | model_scores = [] 22 | for model_name in os.listdir(dir_path): 23 | score = re.search(r'-(\d+(?:\.\d+)?).pth', model_name) 24 | if score is not None: 25 | score = score.group(0)[1:-4] 26 | model_scores.append((model_name, score)) 27 | model_score = sorted(model_scores, key=lambda x: x[1]) 28 | best_model_name = model_score[-1][0] 29 | best_model_path = os.path.join(dir_path, best_model_name) 30 | return best_model_path 31 | 32 | 33 | BASE_EXPERIMENT_NAME = 'fpn-lovasz-se-resnext50-006' 34 | EXPERIMENT_NAME = 'fpn-lovasz-se-resnext50-006-after-001' 35 | N_FOLDS = 5 36 | FOLDS = list(range(N_FOLDS)) 37 | BATCH_SIZE = 16 38 | IMAGE_SIZE = (128, 128) 39 | OUTPUT_SIZE = (101, 101) 40 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148.csv' 41 | LR = 0.005 42 | SAVE_DIR = f'/workdir/data/experiments/{EXPERIMENT_NAME}' 43 | 44 | 45 | class CosineAnnealingLR: 46 | def __init__(self, base_lr, T_max, eta_min=0.): 47 | self.T_max = T_max 48 | self.eta_min = eta_min 49 | self.base_lr = base_lr 50 | 51 | def __call__(self, epoch): 52 | return self.eta_min + (self.base_lr - self.eta_min) \ 53 | * (1 + math.cos(math.pi * (epoch % self.T_max) / self.T_max)) / 2 54 | 55 | 56 | cos_ann = CosineAnnealingLR(LR, 50, eta_min=LR*0.1) 57 | 58 | 59 | @argus.callbacks.on_epoch_start 60 | def update_lr(state: argus.engine.State): 61 | lr = cos_ann(state.epoch) 62 | state.model.set_lr(lr) 63 | state.logger.info(f"Set lr: {lr}") 64 | 65 | 66 | def train_fold(save_dir, train_folds, val_folds, model_path): 67 | depth_trns = SimpleDepthTransform() 68 | train_trns = SaltTransform(IMAGE_SIZE, True, 'crop') 69 | val_trns = SaltTransform(IMAGE_SIZE, False, 'crop') 70 | train_dataset = SaltDataset(TRAIN_FOLDS_PATH, train_folds, train_trns, depth_trns) 71 | val_dataset = SaltDataset(TRAIN_FOLDS_PATH, val_folds, val_trns, depth_trns) 72 | train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, 73 | drop_last=True, num_workers=4) 74 | val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4) 75 | 76 | model = load_model(model_path) 77 | model.nn_module.final = None 78 | model.prediction_transform.segm_thresh = 0 79 | model.loss = LovaszProbLoss(lovasz_weight=0.5, prob_weight=0.5) 80 | 81 | callbacks = [ 82 | MonitorCheckpoint(save_dir, monitor='val_crop_iout', max_saves=3, copy_last=False), 83 | LoggingToFile(os.path.join(save_dir, 'log.txt')), 84 | update_lr 85 | ] 86 | 87 | model.fit(train_loader, 88 | val_loader=val_loader, 89 | max_epochs=300, 90 | callbacks=callbacks, 91 | metrics=['crop_iout']) 92 | 93 | 94 | if __name__ == "__main__": 95 | if not os.path.exists(SAVE_DIR): 96 | os.makedirs(SAVE_DIR) 97 | else: 98 | print(f"Folder {SAVE_DIR} already exists.") 99 | 100 | with open(os.path.join(SAVE_DIR, 'source.py'), 'w') as outfile: 101 | outfile.write(open(__file__).read()) 102 | 103 | for i in range(len(FOLDS)): 104 | val_folds = [FOLDS[i]] 105 | train_folds = FOLDS[:i] + FOLDS[i + 1:] 106 | save_fold_dir = os.path.join(SAVE_DIR, f'fold_{FOLDS[i]}') 107 | print(f"Val folds: {val_folds}, Train folds: {train_folds}") 108 | print(f"Fold save dir {save_fold_dir}") 109 | model_path = get_best_model_path(join('/workdir/data/experiments', BASE_EXPERIMENT_NAME, 'fold_%d' % i)) 110 | print(f'Base model path: {model_path}') 111 | train_fold(save_fold_dir, train_folds, val_folds, model_path) 112 | -------------------------------------------------------------------------------- /notebooks/cv_check.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import numpy as np \n", 10 | "import pandas as pd\n", 11 | "import cv2\n", 12 | "\n", 13 | "import os\n", 14 | "from os.path import join\n", 15 | "\n", 16 | "from src import config\n", 17 | "from src.metrics import compute_eval_metric\n", 18 | "\n", 19 | "from src.transforms import CenterCrop\n", 20 | "\n", 21 | "import matplotlib.pyplot as plt\n", 22 | "%matplotlib inline" 23 | ] 24 | }, 25 | { 26 | "cell_type": "code", 27 | "execution_count": null, 28 | "metadata": {}, 29 | "outputs": [], 30 | "source": [ 31 | "PREDICTION_DIR = '/workdir/data/predictions/mos-fpn-lovasz-se-resnext50-001'\n", 32 | "TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148_mos_emb_1.csv'\n", 33 | "SEGM_THRESH = 0.5\n", 34 | "PROB_THRESH = 0.6" 35 | ] 36 | }, 37 | { 38 | "cell_type": "code", 39 | "execution_count": null, 40 | "metadata": { 41 | "scrolled": false 42 | }, 43 | "outputs": [], 44 | "source": [ 45 | "folds_df = pd.read_csv(TRAIN_FOLDS_PATH)\n", 46 | "score_lst = []\n", 47 | "probs_df_dict = dict()\n", 48 | "\n", 49 | "crop = CenterCrop((101, 101))\n", 50 | "\n", 51 | "for i, row in folds_df.iterrows():\n", 52 | " if row.fold == 5:\n", 53 | " continue\n", 54 | " \n", 55 | " true_mask = cv2.imread(row.mask_path, cv2.IMREAD_GRAYSCALE)\n", 56 | " pred_mask_path = join(PREDICTION_DIR, f'fold_{row.fold}', 'val', row.id+'.png')\n", 57 | " \n", 58 | " if row.fold not in probs_df_dict:\n", 59 | " probs_path = join(PREDICTION_DIR, f'fold_{row.fold}', 'val', 'probs.csv')\n", 60 | " probs_df_dict[row.fold] = pd.read_csv(probs_path, index_col='id')\n", 61 | " \n", 62 | " prob = probs_df_dict[row.fold].loc[row.id].prob\n", 63 | " prob = prob > PROB_THRESH\n", 64 | " \n", 65 | " prob_mask = cv2.imread(pred_mask_path, cv2.IMREAD_GRAYSCALE)\n", 66 | " pred_mask = (prob_mask / 255.0) > SEGM_THRESH\n", 67 | " pred_mask *= prob\n", 68 | " true_mask = true_mask.astype(bool).astype(np.uint8)\n", 69 | " pred_mask = pred_mask.astype(bool).astype(np.uint8)\n", 70 | " \n", 71 | " score = compute_eval_metric(crop(true_mask), pred_mask)\n", 72 | " score_lst.append((row.id, score))\n", 73 | " \n", 74 | "# if score < 0.2:\n", 75 | "# print(score, row.id, probs_df_dict[row.fold].loc[row.id].prob)\n", 76 | "# image_path = join(config.TRAIN_DIR, 'images', row.id+'.png')\n", 77 | "# image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)\n", 78 | " \n", 79 | "# f, axarr = plt.subplots(1, 4)\n", 80 | "# f.set_figheight(4)\n", 81 | "# f.set_figwidth(12)\n", 82 | "# axarr[0].imshow(image)\n", 83 | "# axarr[0].set_title('image')\n", 84 | "# axarr[1].imshow(true_mask)\n", 85 | "# axarr[1].set_title('true')\n", 86 | "# axarr[2].imshow(pred_mask)\n", 87 | "# axarr[2].set_title('pred')\n", 88 | "# axarr[3].imshow(prob_mask)\n", 89 | "# axarr[3].set_title('prob mask')\n", 90 | "# plt.show()\n" 91 | ] 92 | }, 93 | { 94 | "cell_type": "code", 95 | "execution_count": null, 96 | "metadata": {}, 97 | "outputs": [], 98 | "source": [ 99 | "print(np.mean([score for id, score in score_lst]))\n", 100 | "plt.hist([score for id, score in score_lst], bins=20)" 101 | ] 102 | }, 103 | { 104 | "cell_type": "code", 105 | "execution_count": null, 106 | "metadata": {}, 107 | "outputs": [], 108 | "source": [] 109 | } 110 | ], 111 | "metadata": { 112 | "kernelspec": { 113 | "display_name": "Python 3", 114 | "language": "python", 115 | "name": "python3" 116 | }, 117 | "language_info": { 118 | "codemirror_mode": { 119 | "name": "ipython", 120 | "version": 3 121 | }, 122 | "file_extension": ".py", 123 | "mimetype": "text/x-python", 124 | "name": "python", 125 | "nbconvert_exporter": "python", 126 | "pygments_lexer": "ipython3", 127 | "version": "3.6.5" 128 | } 129 | }, 130 | "nbformat": 4, 131 | "nbformat_minor": 2 132 | } 133 | -------------------------------------------------------------------------------- /unused/train_mean_teacher_folds.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | 4 | import argus 5 | from argus.callbacks import MonitorCheckpoint, EarlyStopping, LoggingToFile 6 | 7 | from torch.utils.data import DataLoader 8 | 9 | from src.dataset import SaltDataset, SaltTestDataset 10 | from src.transforms import SimpleDepthTransform, SaltTransform 11 | from src.lr_scheduler import ReduceLROnPlateau 12 | from src.argus_models import SaltMeanTeacherModel 13 | from src import config 14 | 15 | from src.nick_zoo.resnet_blocks import resnet152 16 | 17 | 18 | EXPERIMENT_NAME = 'flex-fpn-mt-resnet152-001' 19 | TRAIN_BATCH_SIZE = 10 20 | VAL_BATCH_SIZE = 10 21 | UNLABELED_BATCH = 6 22 | IMAGE_SIZE = (128, 128) 23 | OUTPUT_SIZE = (101, 101) 24 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148.csv' 25 | TEST_DIR = '/workdir/data/test/images148' 26 | SAVE_DIR = f'/workdir/data/experiments/{EXPERIMENT_NAME}' 27 | FOLDS = list(range(config.N_FOLDS)) 28 | PARAMS = { 29 | 'nn_module': ('UNetFPNFlexProb', { 30 | 'num_classes': 1, 31 | 'num_channels': 3, 32 | 'blocks': resnet152, 33 | 'final': 'sigmoid', 34 | 'dropout_2d': 0.2, 35 | 'is_deconv': True, 36 | 'deflation': 4, 37 | 'use_first_pool': False, 38 | 'skip_dropout': True, 39 | 'pretrain': 'resnet152', 40 | 'pretrain_layers': [True for _ in range(5)], 41 | 'fpn_layers': [16, 32, 64, 128] 42 | }), 43 | 'loss': ('FbBceProbLoss', { 44 | 'fb_weight': 0.95, 45 | 'fb_beta': 2, 46 | 'bce_weight': 0.9, 47 | 'prob_weight': 0.85 48 | }), 49 | 'prediction_transform': ('ProbOutputTransform', { 50 | 'segm_thresh': 0.5, 51 | 'prob_thresh': 0.5, 52 | }), 53 | 'mean_teacher': { 54 | 'alpha': 0.99, 55 | 'rampup_length': 10, 56 | 'unlabeled_batch': UNLABELED_BATCH, 57 | 'consistency_segm_weight': 0.3, 58 | 'consistency_prob_weight': 0.3 59 | }, 60 | 'optimizer': ('Adam', {'lr': 0.0001}), 61 | 'device': 'cuda' 62 | } 63 | 64 | 65 | @argus.callbacks.on_epoch_start 66 | def update_model_epoch(state: argus.engine.State): 67 | state.model.epoch = state.epoch 68 | 69 | 70 | def train_fold(save_dir, train_folds, val_folds): 71 | depth_trns = SimpleDepthTransform() 72 | train_trns = SaltTransform(IMAGE_SIZE, True, 'crop') 73 | val_trns = SaltTransform(IMAGE_SIZE, False, 'crop') 74 | train_dataset = SaltDataset(TRAIN_FOLDS_PATH, train_folds, train_trns, depth_trns) 75 | val_dataset = SaltDataset(TRAIN_FOLDS_PATH, val_folds, val_trns, depth_trns) 76 | train_loader = DataLoader(train_dataset, batch_size=TRAIN_BATCH_SIZE, shuffle=True, 77 | drop_last=True, num_workers=4) 78 | val_loader = DataLoader(val_dataset, batch_size=VAL_BATCH_SIZE, shuffle=False, num_workers=4) 79 | test_dataset = SaltTestDataset(test_dir=TEST_DIR, transform=val_trns, depth_transform=depth_trns) 80 | 81 | model = SaltMeanTeacherModel(PARAMS) 82 | model.test_dataset = test_dataset 83 | 84 | callbacks = [ 85 | MonitorCheckpoint(save_dir, monitor='val_crop_iout', max_saves=3, copy_last=False), 86 | EarlyStopping(monitor='val_crop_iout', patience=100), 87 | ReduceLROnPlateau(monitor='val_crop_iout', patience=25, factor=0.72, min_lr=1e-8), 88 | LoggingToFile(os.path.join(save_dir, 'log.txt')), 89 | update_model_epoch 90 | ] 91 | 92 | model.fit(train_loader, 93 | val_loader=val_loader, 94 | max_epochs=700, 95 | callbacks=callbacks, 96 | metrics=['crop_iout']) 97 | 98 | 99 | if __name__ == "__main__": 100 | if not os.path.exists(SAVE_DIR): 101 | os.makedirs(SAVE_DIR) 102 | else: 103 | print(f"Folder {SAVE_DIR} already exists.") 104 | 105 | with open(os.path.join(SAVE_DIR, 'source.py'), 'w') as outfile: 106 | outfile.write(open(__file__).read()) 107 | 108 | with open(os.path.join(SAVE_DIR, 'params.json'), 'w') as outfile: 109 | json.dump(PARAMS, outfile) 110 | 111 | for i in range(len(FOLDS)): 112 | val_folds = [FOLDS[i]] 113 | train_folds = FOLDS[:i] + FOLDS[i + 1:] 114 | save_fold_dir = os.path.join(SAVE_DIR, f'fold_{FOLDS[i]}') 115 | print(f"Val folds: {val_folds}, Train folds: {train_folds}") 116 | print(f"Fold save dir {save_fold_dir}") 117 | train_fold(save_fold_dir, train_folds, val_folds) 118 | -------------------------------------------------------------------------------- /src/argus_models.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | from argus import Model 5 | from argus.utils import detach_tensors, to_device 6 | 7 | from src.nn_modules import UNetProbResNet 8 | from src.romul_zoo.senet import SeResnextProb50 9 | from src.nick_zoo.unet_flex import UNetFlexProb 10 | from src.nick_zoo.unet_flex_fpn import UNetFPNFlexProb 11 | from src.nick_zoo.senet_fpn import SeResnextFPNProb50 12 | from src.losses import FbBceProbLoss, ConsistencyLoss, LovaszProbLoss 13 | from src.transforms import ProbOutputTransform 14 | from src.metrics import SaltIOUT, SaltCropIOUT 15 | from src.utils import sigmoid_rampup 16 | 17 | 18 | class SaltProbModel(Model): 19 | nn_module = UNetProbResNet 20 | loss = FbBceProbLoss 21 | prediction_transform = ProbOutputTransform 22 | 23 | 24 | class SaltMetaModel(Model): 25 | nn_module = { 26 | 'UNetProbResNet': UNetProbResNet, 27 | 'SeResnextProb50': SeResnextProb50, 28 | 'UNetFlexProb': UNetFlexProb, 29 | 'UNetFPNFlexProb': UNetFPNFlexProb, 30 | 'SeResnextFPNProb50': SeResnextFPNProb50 31 | } 32 | loss = { 33 | 'FbBceProbLoss': FbBceProbLoss, 34 | 'LovaszProbLoss': LovaszProbLoss 35 | } 36 | prediction_transform = { 37 | 'ProbOutputTransform': ProbOutputTransform 38 | } 39 | 40 | 41 | class SaltMeanTeacherModel(Model): 42 | nn_module = { 43 | 'UNetProbResNet': UNetProbResNet, 44 | 'SeResnextProb50': SeResnextProb50, 45 | 'UNetFlexProb': UNetFlexProb, 46 | 'UNetFPNFlexProb': UNetFPNFlexProb, 47 | 'SeResnextFPNProb50': SeResnextFPNProb50 48 | } 49 | loss = { 50 | 'FbBceProbLoss': FbBceProbLoss, 51 | 'LovaszProbLoss': LovaszProbLoss 52 | } 53 | prediction_transform = { 54 | 'ProbOutputTransform': ProbOutputTransform 55 | } 56 | 57 | def __init__(self, params): 58 | super().__init__(params) 59 | self.alpha = params['mean_teacher']['alpha'] 60 | self.rampup_length = params['mean_teacher']['rampup_length'] 61 | self.unlabeled_batch = params['mean_teacher']['unlabeled_batch'] 62 | self.test_dataset = None 63 | self.consistency_loss = ConsistencyLoss( 64 | params['mean_teacher']['consistency_segm_weight'], 65 | params['mean_teacher']['consistency_prob_weight'], 66 | ) 67 | self.epoch = 0 68 | self.teacher = self._build_nn_module(self.params) 69 | self.teacher.to(self.device) 70 | self.teacher.train() 71 | for param in self.teacher.parameters(): 72 | param.detach_() 73 | 74 | def update_teacher(self): 75 | for t_param, s_param in zip(self.teacher.parameters(), 76 | self.nn_module.parameters()): 77 | t_param.data.mul_(self.alpha).add_(1 - self.alpha, s_param.data) 78 | 79 | def sample_unlabeled_input(self): 80 | assert self.test_dataset is not None 81 | indices = np.random.randint(0, len(self.test_dataset), size=self.unlabeled_batch) 82 | samples = [self.test_dataset[idx] for idx in indices] 83 | return torch.stack(samples, dim=0) 84 | 85 | def prepare_unlabeled_batch(self, batch, device): 86 | input, trg = batch 87 | unlabeled_input = self.sample_unlabeled_input() 88 | input = torch.cat([input, unlabeled_input], dim=0) 89 | return to_device(input, device), to_device(trg, device) 90 | 91 | def train_step(self, batch) -> dict: 92 | if not self.nn_module.training: 93 | self.nn_module.train() 94 | self.optimizer.zero_grad() 95 | 96 | input, target = self.prepare_unlabeled_batch(batch, self.device) 97 | student_pred = self.nn_module(input) 98 | with torch.no_grad(): 99 | teacher_pred = self.teacher(input) 100 | 101 | consistency_weight = sigmoid_rampup(self.epoch, self.rampup_length) 102 | loss = consistency_weight * self.consistency_loss(student_pred, teacher_pred) 103 | student_pred = [pred[:target.size(0)] for pred in student_pred] 104 | 105 | loss += self.loss(student_pred, target) 106 | loss.backward() 107 | self.optimizer.step() 108 | self.update_teacher() 109 | prediction = detach_tensors(student_pred) 110 | target = detach_tensors(target) 111 | return { 112 | 'prediction': self.prediction_transform(prediction), 113 | 'target': target, 114 | 'loss': loss.item() 115 | } 116 | -------------------------------------------------------------------------------- /notebooks/mean_submission.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import os\n", 10 | "from os.path import join\n", 11 | "import numpy as np\n", 12 | "import pandas as pd\n", 13 | "import cv2\n", 14 | "\n", 15 | "from src import config\n", 16 | "from src.utils import RLenc, make_dir\n", 17 | "\n", 18 | "PREDICTION_DIRS = [\n", 19 | " '/workdir/data/predictions/mos-fpn-lovasz-se-resnext50-001/'\n", 20 | "]\n", 21 | "FOLDS = [0, 1, 2, 3, 4, 5]\n", 22 | "FOLD_DIRS = [join(p, 'fold_%d'%f) for p in PREDICTION_DIRS for f in FOLDS]\n", 23 | "\n", 24 | "PREDICTION_DIRS = [\n", 25 | " '/workdir/data/predictions/fpn-lovasz-se-resnext50-006-after-001/',\n", 26 | "]\n", 27 | "\n", 28 | "FOLDS = [0, 1, 2, 3, 4]\n", 29 | "FOLD_DIRS += [join(p, 'fold_%d'%f) for p in PREDICTION_DIRS for f in FOLDS]\n", 30 | "\n", 31 | "segm_thresh = 0.4\n", 32 | "prob_thresh = 0.5\n", 33 | "\n", 34 | "SAVE_NAME = 'mean-005-0.4'\n", 35 | "\n", 36 | "make_dir(f'/workdir/data/{SAVE_NAME}')" 37 | ] 38 | }, 39 | { 40 | "cell_type": "code", 41 | "execution_count": null, 42 | "metadata": {}, 43 | "outputs": [], 44 | "source": [ 45 | "FOLD_DIRS" 46 | ] 47 | }, 48 | { 49 | "cell_type": "code", 50 | "execution_count": null, 51 | "metadata": {}, 52 | "outputs": [], 53 | "source": [ 54 | "def get_mean_probs_df():\n", 55 | " probs_df_lst = []\n", 56 | " for fold_dir in FOLD_DIRS:\n", 57 | " probs_df = pd.read_csv(join(fold_dir, 'test', 'probs.csv'), index_col='id')\n", 58 | " probs_df_lst.append(probs_df)\n", 59 | "\n", 60 | " mean_probs_df = probs_df_lst[0].copy()\n", 61 | " for probs_df in probs_df_lst[1:]:\n", 62 | " mean_probs_df.prob += probs_df.prob\n", 63 | " mean_probs_df.prob /= len(probs_df_lst)\n", 64 | "\n", 65 | " return mean_probs_df" 66 | ] 67 | }, 68 | { 69 | "cell_type": "code", 70 | "execution_count": null, 71 | "metadata": {}, 72 | "outputs": [], 73 | "source": [ 74 | "mean_probs_df = get_mean_probs_df()" 75 | ] 76 | }, 77 | { 78 | "cell_type": "code", 79 | "execution_count": null, 80 | "metadata": {}, 81 | "outputs": [], 82 | "source": [ 83 | "sample_submition = pd.read_csv(config.SAMPLE_SUBM_PATH)\n", 84 | "\n", 85 | "for i, row in sample_submition.iterrows():\n", 86 | " pred_name = row.id+'.png'\n", 87 | " pred_lst = []\n", 88 | " for fold_dir in FOLD_DIRS:\n", 89 | " pred_path = join(fold_dir, 'test', pred_name)\n", 90 | " pred = cv2.imread(pred_path, cv2.IMREAD_GRAYSCALE)\n", 91 | " pred = pred / 255\n", 92 | " pred_lst.append(pred)\n", 93 | " \n", 94 | " mean_pred = np.mean(pred_lst, axis=0)\n", 95 | " prob = mean_probs_df.loc[row.id].prob\n", 96 | " \n", 97 | " pred = mean_pred > segm_thresh\n", 98 | " prob = int(prob > prob_thresh)\n", 99 | " pred = (pred * prob).astype(np.uint8)\n", 100 | " \n", 101 | " if np.all(pred == 1):\n", 102 | " pred[:] = 0\n", 103 | " print('Full mask to empty', pred_name)\n", 104 | "\n", 105 | " rle_mask = RLenc(pred)\n", 106 | " cv2.imwrite(f'/workdir/data/{SAVE_NAME}/{pred_name}', pred * 255)\n", 107 | " row.rle_mask = rle_mask\n", 108 | "\n", 109 | "sample_submition.to_csv(f'/workdir/data/submissions/{SAVE_NAME}.csv', index=False) " 110 | ] 111 | }, 112 | { 113 | "cell_type": "code", 114 | "execution_count": null, 115 | "metadata": {}, 116 | "outputs": [], 117 | "source": [ 118 | "import matplotlib.pyplot as plt\n", 119 | "%matplotlib inline\n", 120 | "\n", 121 | "plt.hist(mean_probs_df.prob.values, bins=20)" 122 | ] 123 | }, 124 | { 125 | "cell_type": "code", 126 | "execution_count": null, 127 | "metadata": {}, 128 | "outputs": [], 129 | "source": [] 130 | } 131 | ], 132 | "metadata": { 133 | "kernelspec": { 134 | "display_name": "Python 3", 135 | "language": "python", 136 | "name": "python3" 137 | }, 138 | "language_info": { 139 | "codemirror_mode": { 140 | "name": "ipython", 141 | "version": 3 142 | }, 143 | "file_extension": ".py", 144 | "mimetype": "text/x-python", 145 | "name": "python", 146 | "nbconvert_exporter": "python", 147 | "pygments_lexer": "ipython3", 148 | "version": "3.6.5" 149 | } 150 | }, 151 | "nbformat": 4, 152 | "nbformat_minor": 2 153 | } 154 | -------------------------------------------------------------------------------- /notebooks/make_submission.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import os\n", 10 | "from os.path import join\n", 11 | "import numpy as np\n", 12 | "import pandas as pd\n", 13 | "import cv2\n", 14 | "\n", 15 | "from src import config\n", 16 | "from src.utils import RLenc\n", 17 | "\n", 18 | "PREDICTION_DIR = '/workdir/data/predictions/train-folds-005'\n", 19 | "FOLDS = [0, 1, 2, 3, 4]\n", 20 | "\n", 21 | "\n", 22 | "segm_thresh = 0.5\n", 23 | "prob_thresh = 0.5" 24 | ] 25 | }, 26 | { 27 | "cell_type": "code", 28 | "execution_count": null, 29 | "metadata": {}, 30 | "outputs": [], 31 | "source": [ 32 | "def get_mean_probs_df(pred_dir):\n", 33 | " probs_df_lst = []\n", 34 | " for i in range(len(FOLDS)):\n", 35 | " fold_dir = os.path.join(pred_dir, f'fold_{FOLDS[i]}')\n", 36 | " probs_df = pd.read_csv(join(fold_dir, 'test', 'probs.csv'), index_col='id')\n", 37 | " probs_df_lst.append(probs_df)\n", 38 | "\n", 39 | " mean_probs_df = probs_df_lst[0].copy()\n", 40 | " for probs_df in probs_df_lst[1:]:\n", 41 | " mean_probs_df.prob += probs_df.prob\n", 42 | " mean_probs_df.prob /= len(probs_df_lst)\n", 43 | "\n", 44 | " return mean_probs_df" 45 | ] 46 | }, 47 | { 48 | "cell_type": "code", 49 | "execution_count": null, 50 | "metadata": {}, 51 | "outputs": [], 52 | "source": [ 53 | "mean_probs_df = get_mean_probs_df(PREDICTION_DIR)" 54 | ] 55 | }, 56 | { 57 | "cell_type": "code", 58 | "execution_count": null, 59 | "metadata": {}, 60 | "outputs": [], 61 | "source": [ 62 | "sample_submition = pd.read_csv(config.SAMPLE_SUBM_PATH)\n", 63 | "\n", 64 | "for i, row in sample_submition.iterrows():\n", 65 | " pred_name = row.id+'.npy'\n", 66 | " pred_lst = []\n", 67 | " for i in range(len(FOLDS)):\n", 68 | " fold_dir = os.path.join(PREDICTION_DIR, f'fold_{FOLDS[i]}')\n", 69 | " pred_path = join(fold_dir, 'test', pred_name)\n", 70 | " pred = np.load(pred_path)\n", 71 | " pred_lst.append(pred)\n", 72 | " \n", 73 | " mean_pred = np.mean(pred_lst, axis=0)\n", 74 | " prob = mean_probs_df.loc[row.id].prob\n", 75 | " \n", 76 | " pred = mean_pred > segm_thresh\n", 77 | " prob = int(prob > prob_thresh)\n", 78 | " pred = (pred * prob).astype(np.uint8)\n", 79 | "\n", 80 | " rle_mask = RLenc(pred)\n", 81 | " row.rle_mask = rle_mask\n", 82 | "\n", 83 | "sample_submition.to_csv('/workdir/data/submissions/train-folds-005-0.45-mean.csv', index=False) " 84 | ] 85 | }, 86 | { 87 | "cell_type": "code", 88 | "execution_count": null, 89 | "metadata": {}, 90 | "outputs": [], 91 | "source": [ 92 | "import matplotlib.pyplot as plt\n", 93 | "%matplotlib inline\n", 94 | "\n", 95 | "plt.hist(mean_probs_df.prob.values, bins=20)" 96 | ] 97 | }, 98 | { 99 | "cell_type": "code", 100 | "execution_count": null, 101 | "metadata": {}, 102 | "outputs": [], 103 | "source": [ 104 | "sample_submition = pd.read_csv(config.SAMPLE_SUBM_PATH)\n", 105 | "input_dir = '/workdir/data/predictions/mean-005-0.40_draw/'\n", 106 | "for i, row in sample_submition.iterrows():\n", 107 | " pred_name = row.id+'.png'\n", 108 | " pred = cv2.imread(join(input_dir, pred_name), cv2.IMREAD_GRAYSCALE)\n", 109 | " pred = pred > 0\n", 110 | " rle_mask = RLenc(pred.astype(np.uint8))\n", 111 | " row.rle_mask = rle_mask\n", 112 | "\n", 113 | "sample_submition.to_csv('/workdir/data/submissions/mean-005-0.40_draw.csv', index=False) " 114 | ] 115 | }, 116 | { 117 | "cell_type": "code", 118 | "execution_count": null, 119 | "metadata": {}, 120 | "outputs": [], 121 | "source": [ 122 | "sample_submition.to_csv('/workdir/data/submissions/mean-003-0.40_draw.csv', index=False) " 123 | ] 124 | }, 125 | { 126 | "cell_type": "code", 127 | "execution_count": null, 128 | "metadata": {}, 129 | "outputs": [], 130 | "source": [] 131 | } 132 | ], 133 | "metadata": { 134 | "kernelspec": { 135 | "display_name": "Python 3", 136 | "language": "python", 137 | "name": "python3" 138 | }, 139 | "language_info": { 140 | "codemirror_mode": { 141 | "name": "ipython", 142 | "version": 3 143 | }, 144 | "file_extension": ".py", 145 | "mimetype": "text/x-python", 146 | "name": "python", 147 | "nbconvert_exporter": "python", 148 | "pygments_lexer": "ipython3", 149 | "version": "3.6.5" 150 | } 151 | }, 152 | "nbformat": 4, 153 | "nbformat_minor": 2 154 | } 155 | -------------------------------------------------------------------------------- /mosaic/mosaic_api_examples.py.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import json \n", 10 | "import numpy as np\n", 11 | "import cv2\n", 12 | "import os\n", 13 | "from os.path import join\n", 14 | "from skimage.restoration import inpaint_biharmonic\n", 15 | "\n", 16 | "from mosaic.mosaic_api import SaltData, Mosaic, Mosaics\n", 17 | "from src.transforms import CenterCrop\n", 18 | "\n", 19 | "import matplotlib.pyplot as plt\n", 20 | "%matplotlib inline" 21 | ] 22 | }, 23 | { 24 | "cell_type": "code", 25 | "execution_count": null, 26 | "metadata": {}, 27 | "outputs": [], 28 | "source": [ 29 | "def imshow(image, figsize=(3, 3)):\n", 30 | " plt.figure(figsize=figsize)\n", 31 | " plt.imshow(image)\n", 32 | " plt.show()" 33 | ] 34 | }, 35 | { 36 | "cell_type": "code", 37 | "execution_count": null, 38 | "metadata": {}, 39 | "outputs": [], 40 | "source": [ 41 | "mosaic_path = '/workdir/data/mosaic/pazzles_6013.csv'\n", 42 | "saltdata = SaltData(mosaic_csv_path=mosaic_path)" 43 | ] 44 | }, 45 | { 46 | "cell_type": "code", 47 | "execution_count": null, 48 | "metadata": {}, 49 | "outputs": [], 50 | "source": [ 51 | "len(saltdata.train_ids), len(saltdata.test_ids), len(saltdata.mosaics.mosaic_id2mosaic)" 52 | ] 53 | }, 54 | { 55 | "cell_type": "code", 56 | "execution_count": null, 57 | "metadata": {}, 58 | "outputs": [], 59 | "source": [ 60 | "# Mosaic list \n", 61 | "mosaic = saltdata.mosaics[79]\n", 62 | "mosaic" 63 | ] 64 | }, 65 | { 66 | "cell_type": "code", 67 | "execution_count": null, 68 | "metadata": {}, 69 | "outputs": [], 70 | "source": [ 71 | "# Get neighbors of tile with id \n", 72 | "neighbors = saltdata.get_neighbors('a37249665e')\n", 73 | "neighbors" 74 | ] 75 | }, 76 | { 77 | "cell_type": "code", 78 | "execution_count": null, 79 | "metadata": {}, 80 | "outputs": [], 81 | "source": [ 82 | "# Neighbors visualization \n", 83 | "vis_img = saltdata.draw_visualization(neighbors)\n", 84 | "imshow(vis_img)" 85 | ] 86 | }, 87 | { 88 | "cell_type": "code", 89 | "execution_count": null, 90 | "metadata": {}, 91 | "outputs": [], 92 | "source": [ 93 | "# Full mosaic visualization \n", 94 | "vis_img = saltdata.draw_visualization(mosaic.array)\n", 95 | "imshow(vis_img, figsize=(7, 7))\n", 96 | "mosaic" 97 | ] 98 | }, 99 | { 100 | "cell_type": "code", 101 | "execution_count": null, 102 | "metadata": {}, 103 | "outputs": [], 104 | "source": [ 105 | "print('c2ec2c9de4' in saltdata.mosaics)\n", 106 | "print(saltdata.in_mosaics('c2ec2c9de4'))\n", 107 | "print('c2ec2c9de4' in mosaic)\n", 108 | "print(saltdata.in_train('c2ec2c9de4'))\n", 109 | "print(saltdata.mosaics.id2mosaic_id['c2ec2c9de4'])\n", 110 | "saltdata.mosaics.id2mosaic['c2ec2c9de4']" 111 | ] 112 | }, 113 | { 114 | "cell_type": "code", 115 | "execution_count": null, 116 | "metadata": {}, 117 | "outputs": [], 118 | "source": [ 119 | "center = CenterCrop((148, 148))" 120 | ] 121 | }, 122 | { 123 | "cell_type": "code", 124 | "execution_count": null, 125 | "metadata": {}, 126 | "outputs": [], 127 | "source": [ 128 | "stacked_image, image_unknown_mask = saltdata.get_stacked_images(neighbors, return_unknown_mask=True)\n", 129 | "inpaint_image = inpaint_biharmonic(center(stacked_image), center(image_unknown_mask))\n", 130 | "\n", 131 | "imshow(stacked_image)\n", 132 | "imshow(image_unknown_mask)\n", 133 | "imshow(inpaint_image)" 134 | ] 135 | }, 136 | { 137 | "cell_type": "code", 138 | "execution_count": null, 139 | "metadata": {}, 140 | "outputs": [], 141 | "source": [ 142 | "stacked_mask, mask_unknown_mask = saltdata.get_stacked_masks(neighbors, return_unknown_mask=True)\n", 143 | "if np.sum(center(mask_unknown_mask)):\n", 144 | " inpaint_mask = inpaint_biharmonic(center(stacked_mask), center(mask_unknown_mask))\n", 145 | " imshow(inpaint_mask)\n", 146 | "\n", 147 | "imshow(stacked_mask)\n", 148 | "imshow(mask_unknown_mask)" 149 | ] 150 | }, 151 | { 152 | "cell_type": "code", 153 | "execution_count": null, 154 | "metadata": {}, 155 | "outputs": [], 156 | "source": [] 157 | } 158 | ], 159 | "metadata": { 160 | "kernelspec": { 161 | "display_name": "Python 3", 162 | "language": "python", 163 | "name": "python3" 164 | }, 165 | "language_info": { 166 | "codemirror_mode": { 167 | "name": "ipython", 168 | "version": 3 169 | }, 170 | "file_extension": ".py", 171 | "mimetype": "text/x-python", 172 | "name": "python", 173 | "nbconvert_exporter": "python", 174 | "pygments_lexer": "ipython3", 175 | "version": "3.6.5" 176 | } 177 | }, 178 | "nbformat": 4, 179 | "nbformat_minor": 2 180 | } 181 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Argus solution TGS Salt Identification Challenge 2 | 3 | Source code of 14th place solution for [TGS Salt Identification Challenge](https://www.kaggle.com/c/tgs-salt-identification-challenge) by Argus team ([Ruslan Baikulov](https://www.kaggle.com/romul0212), [Nikolay Falaleev](https://www.kaggle.com/nikolasent)). 4 | 5 | 6 | ## Solution 7 | 8 | We used PyTorch 0.4.1 with framework [Argus](https://github.com/lRomul/argus) simplifies the experiments with [different architectures](https://github.com/lRomul/argus-tgs-salt-dev/blob/master/src/argus_models.py) and allows to focus on deep learning trials rather than coding neural networks training and testing scripts. 9 | 10 | ### Data preprocessing 11 | 12 | The original images with size 101x101 px padded to 148x148 px with biharmonic inpaint from `skimage` package. This “padding” performed better for us than reflection or zero padding. 13 | Random crop to the input size 128x128 px, flip in the left-right direction and random linear color augmentation (for brightness and contrast adjustment) were applied. 14 | 15 | ### Model design: 16 | 17 | ![General scheme of the neural networks architectures](readme_img/architecture.png) 18 | 19 | After a series of experiments, we ended with a U-Net like architecture with an SE-ResNeXt50 encoder. Standard decoder blocks enriched with custom-built FPN-style layers. In addition to the segmentation task, an additional classification branch (empty/contains salt tile) added into basic network architecture. 20 | 21 | 22 | ### Models Training 23 | 24 | Loss: Lovasz hinge loss with [elu + 1](https://github.com/lRomul/argus-tgs-salt-dev/blob/e337c9874d3dc5a3274c5426451b314b9d837a46/src/lovasz.py#L110) 25 | Optimizer: SGD with LR 0.01, momentum 0.9, weight_decay 0.0001 26 | Train stages: 27 | 1. EarlyStopping with patience 100; ReduceLROnPlateau with patience=30, factor=0.64, min_lr=1e-8; Lovasz * 0.75 + BCE empty * 0.25. 28 | 2. Cosine annealing learning rate 300 epochs, 50 per cycle; Lovasz * 0.5 + BCE empty * 0.5. 29 | 30 | ### Post-processing 31 | 32 | Averaged results of two training used in the final submission: 33 | SE-ResNeXt50 on 5 random folds. 34 | SE-ResNeXt50 on 6 mosaic based folds (similar mosaics tiles placed in the same fold) without the second training stage. 35 | 36 | Mosaics-based post-processing. We used the Vicens Gaitan’s [Kernel](https://www.kaggle.com/vicensgaitan/salt-jigsaw-puzzle) but not on a raw input dataset, but on images after histogram matching to an average histogram, which helps us to assembly more tiles into mosaics. In addition to extrapolating tiles with vertical masks from train subset on neighbouring tiles, we performed an automatic detection of small missed corners and inpaint them by a polygon with a smooth-curve boundary. Holes in masks were also filled with OpenCV. 37 | 38 | 39 | ### Results 40 | 41 | ![Sample results image](readme_img/postprocess.png) 42 | 43 | Example of the whole mosaic post-processing. Green/blue - salt/empty regions from the train dataset; red - predicted mask; yellow - inpainted by the post-processing (used in the final submission). 44 | 45 | The results are based on a step-by-step improvement of the pipeline, postprocessing, and fair cross-validation. Finally, results were achieved by carefully selected architectures without heavy ensembles of neural nets and second order models. Reasonable cross-validation with the evaluation metric prevented us from overfitting on the public leaderboard. 46 | 47 | For more details on data pre- and post-processing, as well as conducted experiments with neural nets, check out a blog [post](https://nikolasent.github.io/deeplearning/competitions/2018/10/24/Semantic-Segmentation-of-Seismic-Reflection-Images.html). 48 | 49 | 50 | ## Quick setup and start 51 | 52 | ### Requirements 53 | 54 | * Nvidia drivers, CUDA >= 9, cuDNN >= 7 55 | * [Docker](https://www.docker.com/), [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) 56 | 57 | The provided dockerfile is supplied to build image with cuda support and cudnn. 58 | 59 | 60 | ### Preparations 61 | 62 | * Clone the repo, build docker image. 63 | ```bash 64 | git clone https://github.com/lRomul/argus-tgs-salt.git 65 | cd argus-tgs-salt 66 | make build 67 | ``` 68 | 69 | * Download and extract [dataset](https://www.kaggle.com/c/tgs-salt-identification-challenge/data) 70 | * extract train images and masks into `data/train/` 71 | * extract test images into `data/test/` 72 | 73 | * The folder structure should be: 74 | ``` 75 | argus-tgs-salt 76 | ├── data 77 | │   ├── mosaic 78 | │   ├── test 79 | │   └── train 80 | ├── docker 81 | ├── mosaic 82 | ├── notebooks 83 | ├── pipeline 84 | ├── src 85 | └── unused 86 | ``` 87 | 88 | ### Run 89 | 90 | * Run docker container 91 | ```bash 92 | make run 93 | ``` 94 | 95 | * Start full pipeline with postprocessing 96 | ```bash 97 | ./run_pipeline.sh 98 | ``` 99 | 100 | The final submission file will be at: 101 | ``` 102 | data/predictions/mean-005-0.4/submission.csv 103 | ``` 104 | -------------------------------------------------------------------------------- /src/nn_modules.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import torchvision 5 | 6 | 7 | def conv3x3(in_, out): 8 | return nn.Conv2d(in_, out, 3, padding=1) 9 | 10 | 11 | class ConvRelu(nn.Module): 12 | def __init__(self, in_, out): 13 | super().__init__() 14 | self.conv = conv3x3(in_, out) 15 | self.activation = nn.ReLU(inplace=True) 16 | 17 | def forward(self, x): 18 | x = self.conv(x) 19 | x = self.activation(x) 20 | return x 21 | 22 | 23 | class DecoderBlockV2(nn.Module): 24 | def __init__(self, in_channels, middle_channels, out_channels, is_deconv=True): 25 | super(DecoderBlockV2, self).__init__() 26 | self.in_channels = in_channels 27 | 28 | if is_deconv: 29 | """ 30 | Paramaters for Deconvolution were chosen to avoid artifacts, following 31 | link https://distill.pub/2016/deconv-checkerboard/ 32 | """ 33 | 34 | self.block = nn.Sequential( 35 | ConvRelu(in_channels, middle_channels), 36 | nn.ConvTranspose2d(middle_channels, out_channels, kernel_size=4, stride=2, 37 | padding=1), 38 | nn.ReLU(inplace=True) 39 | ) 40 | else: 41 | self.block = nn.Sequential( 42 | nn.Upsample(scale_factor=2, mode='bilinear'), 43 | ConvRelu(in_channels, middle_channels), 44 | ConvRelu(middle_channels, out_channels), 45 | ) 46 | 47 | def forward(self, x): 48 | return self.block(x) 49 | 50 | 51 | class UNetProbResNet(nn.Module): 52 | def __init__(self, encoder_depth, num_classes, num_filters=32, dropout_2d=0.2, 53 | pretrained=False, is_deconv=False): 54 | super().__init__() 55 | self.num_classes = num_classes 56 | self.dropout_2d = dropout_2d 57 | 58 | if encoder_depth == 34: 59 | self.encoder = torchvision.models.resnet34(pretrained=pretrained) 60 | bottom_channel_nr = 512 61 | elif encoder_depth == 101: 62 | self.encoder = torchvision.models.resnet101(pretrained=pretrained) 63 | bottom_channel_nr = 2048 64 | elif encoder_depth == 152: 65 | self.encoder = torchvision.models.resnet152(pretrained=pretrained) 66 | bottom_channel_nr = 2048 67 | else: 68 | raise NotImplementedError('only 34, 101, 152 version of Resnet are implemented') 69 | 70 | self.pool = nn.MaxPool2d(2, 2) 71 | self.relu = nn.ReLU(inplace=True) 72 | self.conv1 = nn.Sequential(self.encoder.conv1, self.encoder.bn1, self.encoder.relu, self.pool) 73 | self.conv2 = self.encoder.layer1 74 | self.conv3 = self.encoder.layer2 75 | self.conv4 = self.encoder.layer3 76 | self.conv5 = self.encoder.layer4 77 | 78 | self.center = DecoderBlockV2(bottom_channel_nr, num_filters * 8 * 2, num_filters * 8, is_deconv) 79 | self.dec5 = DecoderBlockV2(bottom_channel_nr + num_filters * 8, num_filters * 8 * 2, num_filters * 8, is_deconv) 80 | self.dec4 = DecoderBlockV2(bottom_channel_nr // 2 + num_filters * 8, num_filters * 8 * 2, num_filters * 8, 81 | is_deconv) 82 | self.dec3 = DecoderBlockV2(bottom_channel_nr // 4 + num_filters * 8, num_filters * 4 * 2, num_filters * 2, 83 | is_deconv) 84 | self.dec2 = DecoderBlockV2(bottom_channel_nr // 8 + num_filters * 2, num_filters * 2 * 2, num_filters * 2 * 2, 85 | is_deconv) 86 | self.dec1 = DecoderBlockV2(num_filters * 2 * 2, num_filters * 2 * 2, num_filters, is_deconv) 87 | self.dec0 = ConvRelu(num_filters, num_filters) 88 | self.final_conv = nn.Conv2d(num_filters, num_classes, kernel_size=1) 89 | self.sigmoid = nn.Sigmoid() 90 | 91 | self.prob_avg_pool = nn.AdaptiveAvgPool2d(1) 92 | self.prob_fc = nn.Linear(num_filters * 8, 1) 93 | 94 | def forward(self, x): 95 | conv1 = self.conv1(x) 96 | conv2 = self.conv2(conv1) 97 | conv3 = self.conv3(conv2) 98 | conv4 = self.conv4(conv3) 99 | conv5 = self.conv5(conv4) 100 | 101 | pool = self.pool(conv5) 102 | center = self.center(pool) 103 | 104 | dec5 = self.dec5(torch.cat([center, conv5], 1)) 105 | dec4 = self.dec4(torch.cat([dec5, conv4], 1)) 106 | dec3 = self.dec3(torch.cat([dec4, conv3], 1)) 107 | dec2 = self.dec2(torch.cat([dec3, conv2], 1)) 108 | dec1 = self.dec1(dec2) 109 | dec0 = self.dec0(dec1) 110 | 111 | prob = self.prob_avg_pool(center) 112 | prob = prob.view(prob.size(0), -1) 113 | prob = self.prob_fc(prob) 114 | prob = self.sigmoid(prob) 115 | prob = prob.view(-1) 116 | 117 | x = self.final_conv(F.dropout2d(dec0, p=self.dropout_2d)) 118 | x = self.sigmoid(x) 119 | 120 | return x, prob 121 | -------------------------------------------------------------------------------- /src/nick_zoo/layers.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | # Some layers are from https://github.com/pytorch/vision/tree/master/torchvision/models 5 | 6 | nonlinearity = nn.ReLU 7 | 8 | 9 | def conv3x3(in_planes, out_planes, stride=1, padding=1, bias=False): 10 | """3x3 convolution with padding""" 11 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 12 | padding=padding, bias=bias) 13 | 14 | 15 | class Transition(nn.Sequential): 16 | def __init__(self, num_input_features, num_output_features): 17 | super(Transition, self).__init__() 18 | self.add_module('norm', nn.BatchNorm2d(num_input_features)) 19 | self.add_module('relu', nn.ReLU(inplace=True)) 20 | self.add_module('conv', nn.Conv2d(num_input_features, num_output_features, 21 | kernel_size=1, stride=1, bias=False)) 22 | self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2)) 23 | 24 | 25 | class BasicBlock(nn.Module): 26 | expansion = 1 27 | 28 | def __init__(self, inplanes, planes, stride=1, downsample=None): 29 | super(BasicBlock, self).__init__() 30 | self.conv1 = conv3x3(inplanes, planes, stride) 31 | self.bn1 = nn.BatchNorm2d(planes) 32 | self.relu = nonlinearity(inplace=True) 33 | self.conv2 = conv3x3(planes, planes) 34 | self.bn2 = nn.BatchNorm2d(planes) 35 | self.downsample = downsample 36 | self.stride = stride 37 | 38 | def forward(self, x): 39 | residual = x 40 | out = self.conv1(x) 41 | out = self.bn1(out) 42 | out = self.relu(out) 43 | out = self.conv2(out) 44 | out = self.bn2(out) 45 | if self.downsample is not None: 46 | residual = self.downsample(x) 47 | out += residual 48 | out = self.relu(out) 49 | return out 50 | 51 | 52 | class Bottleneck(nn.Module): 53 | expansion = 4 54 | 55 | def __init__(self, inplanes, planes, stride=1, downsample=None): 56 | super(Bottleneck, self).__init__() 57 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 58 | self.bn1 = nn.BatchNorm2d(planes) 59 | self.conv2 = conv3x3(planes, planes, stride=stride, 60 | padding=1, bias=False) 61 | self.bn2 = nn.BatchNorm2d(planes) 62 | self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False) 63 | self.bn3 = nn.BatchNorm2d(planes * self.expansion) 64 | self.relu = nonlinearity(inplace=True) 65 | self.downsample = downsample 66 | self.stride = stride 67 | 68 | def forward(self, x): 69 | residual = x 70 | out = self.conv1(x) 71 | out = self.bn1(out) 72 | out = self.relu(out) 73 | out = self.conv2(out) 74 | out = self.bn2(out) 75 | out = self.relu(out) 76 | out = self.conv3(out) 77 | out = self.bn3(out) 78 | if self.downsample is not None: 79 | residual = self.downsample(x) 80 | out += residual 81 | out = self.relu(out) 82 | return out 83 | 84 | 85 | class DecoderBlock(nn.Module): 86 | def __init__(self, in_channels, n_filters, is_deconv=True, 87 | deflation=4, xavier=False): 88 | super(DecoderBlock, self).__init__() 89 | 90 | # B, C, H, W -> B, C/4, H, W 91 | assert in_channels % deflation == 0, "Incorrect deflation" 92 | self.conv1 = nn.Conv2d(in_channels, in_channels // deflation, 1) 93 | self.norm1 = nn.BatchNorm2d(in_channels // deflation) 94 | self.relu1 = nonlinearity(inplace=True) 95 | 96 | # B, C/4, H, W -> B, C/4, H, W 97 | if is_deconv: 98 | self.upscale = nn.ConvTranspose2d(in_channels // deflation, in_channels // deflation, 3, 99 | stride=2, padding=1, output_padding=1) 100 | else: 101 | self.upscale = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) 102 | self.norm2 = nn.BatchNorm2d(in_channels // deflation) 103 | self.relu2 = nonlinearity(inplace=True) 104 | 105 | # B, C/4, H, W -> B, C, H, W 106 | self.conv3 = nn.Conv2d(in_channels // deflation, n_filters, 1) 107 | self.norm3 = nn.BatchNorm2d(n_filters) 108 | self.relu3 = nonlinearity(inplace=True) 109 | 110 | if xavier: 111 | for m in self.modules(): 112 | if isinstance(m, nn.Conv2d): 113 | nn.init.xavier_uniform_(m.weight) 114 | 115 | def forward(self, x): 116 | x = self.conv1(x) 117 | x = self.norm1(x) 118 | x = self.relu1(x) 119 | x = self.upscale(x) 120 | x = self.norm2(x) 121 | x = self.relu2(x) 122 | x = self.conv3(x) 123 | x = self.norm3(x) 124 | x = self.relu3(x) 125 | return x 126 | 127 | 128 | class FPNBlock(nn.Module): 129 | def __init__(self, in_channels, out_channels, scale_factor=2, 130 | xavier=False): 131 | super(FPNBlock, self).__init__() 132 | 133 | mid_channels = int(round((in_channels + out_channels) / 2)) 134 | self.conv1 = conv3x3(in_channels, mid_channels) 135 | self.relu1 = nonlinearity(inplace=True) 136 | self.bn1 = nn.BatchNorm2d(mid_channels) 137 | self.conv2 = conv3x3(mid_channels, out_channels) 138 | self.relu2 = nonlinearity(inplace=True) 139 | self.bn2 = nn.BatchNorm2d(out_channels) 140 | if scale_factor > 1: 141 | self.upscale = nn.Upsample(scale_factor=scale_factor, 142 | mode='bilinear', align_corners=True) 143 | else: 144 | self.upscale = lambda x: x 145 | 146 | if xavier: 147 | for m in self.modules(): 148 | if isinstance(m, nn.Conv2d): 149 | nn.init.xavier_uniform_(m.weight) 150 | 151 | def forward(self, x): 152 | x = self.conv1(x) 153 | x = self.relu1(x) 154 | x = self.bn1(x) 155 | x = self.conv2(x) 156 | x = self.relu2(x) 157 | x = self.bn2(x) 158 | x = self.upscale(x) 159 | return x 160 | -------------------------------------------------------------------------------- /pipeline/fpn-lovasz-se-resnext50-006/predict_folds.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join 3 | import re 4 | import cv2 5 | import numpy as np 6 | import pandas as pd 7 | 8 | from argus import load_model 9 | 10 | import torch 11 | 12 | from src.transforms import SimpleDepthTransform, SaltTransform, CenterCrop 13 | from src.argus_models import SaltMetaModel 14 | from src.transforms import HorizontalFlip 15 | from src.utils import RLenc, make_dir 16 | from src import config 17 | 18 | EXPERIMENT_NAME = 'fpn-lovasz-se-resnext50-006' 19 | N_FOLDS = 5 20 | FOLDS = list(range(N_FOLDS)) 21 | ORIG_IMAGE_SIZE = (101, 101) 22 | PRED_IMAGE_SIZE = (128, 128) 23 | TRANSFORM_MODE = 'crop' 24 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148.csv' 25 | IMAGES_NAME = '148' 26 | FOLDS_DIR = f'/workdir/data/experiments/{EXPERIMENT_NAME}' 27 | PREDICTION_DIR = f'/workdir/data/predictions/{EXPERIMENT_NAME}' 28 | make_dir(PREDICTION_DIR) 29 | 30 | SEGM_THRESH = 0.5 31 | PROB_THRESH = 0.5 32 | 33 | 34 | class Predictor: 35 | def __init__(self, model_path): 36 | self.model = load_model(model_path) 37 | self.model.nn_module.final = torch.nn.Sigmoid() 38 | self.model.nn_module.eval() 39 | 40 | self.depth_trns = SimpleDepthTransform() 41 | self.crop_trns = CenterCrop(ORIG_IMAGE_SIZE) 42 | self.trns = SaltTransform(PRED_IMAGE_SIZE, False, TRANSFORM_MODE) 43 | 44 | def __call__(self, image): 45 | tensor = self.depth_trns(image, 0) 46 | tensor = self.trns(tensor) 47 | tensor = tensor.unsqueeze(0).to(self.model.device) 48 | 49 | with torch.no_grad(): 50 | segm, prob = self.model.nn_module(tensor) 51 | segm = segm.cpu().numpy()[0][0] 52 | segm = self.crop_trns(segm) 53 | 54 | segm = (segm * 255).astype(np.uint8) 55 | prob = prob.item() 56 | 57 | return segm, prob 58 | 59 | 60 | def pred_val_fold(model_path, fold): 61 | predictor = Predictor(model_path) 62 | folds_df = pd.read_csv(TRAIN_FOLDS_PATH) 63 | fold_df = folds_df[folds_df.fold == fold] 64 | fold_prediction_dir = join(PREDICTION_DIR, f'fold_{fold}', 'val') 65 | make_dir(fold_prediction_dir) 66 | 67 | prob_dict = {'id': [], 'prob': []} 68 | for i, row in fold_df.iterrows(): 69 | image = cv2.imread(row.image_path, cv2.IMREAD_GRAYSCALE) 70 | segm, prob = predictor(image) 71 | segm_save_path = join(fold_prediction_dir, row.id + '.png') 72 | cv2.imwrite(segm_save_path, segm) 73 | 74 | prob_dict['id'].append(row.id) 75 | prob_dict['prob'].append(prob) 76 | 77 | prob_df = pd.DataFrame(prob_dict) 78 | prob_df.to_csv(join(fold_prediction_dir, 'probs.csv'), index=False) 79 | 80 | 81 | def pred_test_fold(model_path, fold): 82 | predictor = Predictor(model_path) 83 | prob_df = pd.read_csv(config.SAMPLE_SUBM_PATH) 84 | prob_df.rename(columns={'rle_mask': 'prob'}, inplace=True) 85 | 86 | fold_prediction_dir = join(PREDICTION_DIR, f'fold_{fold}', 'test') 87 | make_dir(fold_prediction_dir) 88 | 89 | for i, row in prob_df.iterrows(): 90 | image_path = join(config.TEST_DIR, 'images'+IMAGES_NAME, row.id + '.png') 91 | image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) 92 | segm, prob = predictor(image) 93 | row.prob = prob 94 | segm_save_path = join(fold_prediction_dir, row.id + '.png') 95 | cv2.imwrite(segm_save_path, segm) 96 | 97 | prob_df.to_csv(join(fold_prediction_dir, 'probs.csv'), index=False) 98 | 99 | 100 | def get_mean_probs_df(pred_dir): 101 | probs_df_lst = [] 102 | for i in range(len(FOLDS)): 103 | fold_dir = os.path.join(pred_dir, f'fold_{FOLDS[i]}') 104 | probs_df = pd.read_csv(join(fold_dir, 'test', 'probs.csv'), index_col='id') 105 | probs_df_lst.append(probs_df) 106 | 107 | mean_probs_df = probs_df_lst[0].copy() 108 | for probs_df in probs_df_lst[1:]: 109 | mean_probs_df.prob += probs_df.prob 110 | mean_probs_df.prob /= len(probs_df_lst) 111 | 112 | return mean_probs_df 113 | 114 | 115 | def make_mean_submission(): 116 | mean_probs_df = get_mean_probs_df(PREDICTION_DIR) 117 | 118 | sample_submition = pd.read_csv(config.SAMPLE_SUBM_PATH) 119 | 120 | for i, row in sample_submition.iterrows(): 121 | pred_name = row.id + '.png' 122 | segm_lst = [] 123 | for i in range(len(FOLDS)): 124 | fold_dir = os.path.join(PREDICTION_DIR, f'fold_{FOLDS[i]}') 125 | pred_path = join(fold_dir, 'test', pred_name) 126 | segm = cv2.imread(pred_path, cv2.IMREAD_GRAYSCALE) 127 | segm = segm.astype(np.float32) / 255 128 | segm_lst.append(segm) 129 | 130 | mean_segm = np.mean(segm_lst, axis=0) 131 | prob = mean_probs_df.loc[row.id].prob 132 | 133 | pred = mean_segm > SEGM_THRESH 134 | prob = int(prob > PROB_THRESH) 135 | pred = (pred * prob).astype(np.uint8) 136 | 137 | rle_mask = RLenc(pred) 138 | row.rle_mask = rle_mask 139 | 140 | sample_submition.to_csv(join(PREDICTION_DIR, f'{EXPERIMENT_NAME}-mean-subm.csv'), index=False) 141 | 142 | 143 | def get_best_model_path(dir_path): 144 | model_scores = [] 145 | for model_name in os.listdir(dir_path): 146 | score = re.search(r'-(\d+(?:\.\d+)?).pth', model_name) 147 | if score is not None: 148 | score = score.group(0)[1:-4] 149 | model_scores.append((model_name, score)) 150 | model_score = sorted(model_scores, key=lambda x: x[1]) 151 | best_model_name = model_score[-1][0] 152 | best_model_path = os.path.join(dir_path, best_model_name) 153 | return best_model_path 154 | 155 | 156 | if __name__ == "__main__": 157 | for i in range(len(FOLDS)): 158 | print("Predict fold", FOLDS[i]) 159 | fold_dir = os.path.join(FOLDS_DIR, f'fold_{FOLDS[i]}') 160 | best_model_path = get_best_model_path(fold_dir) 161 | print("Model path", best_model_path) 162 | print("Val predict") 163 | pred_val_fold(best_model_path, FOLDS[i]) 164 | print("Test predict") 165 | pred_test_fold(best_model_path, FOLDS[i]) 166 | 167 | print("Mean submission") 168 | make_mean_submission() 169 | -------------------------------------------------------------------------------- /pipeline/mos-fpn-lovasz-se-resnext50-001/predict_folds.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join 3 | import re 4 | import cv2 5 | import numpy as np 6 | import pandas as pd 7 | 8 | from argus import load_model 9 | 10 | import torch 11 | 12 | from src.transforms import SimpleDepthTransform, SaltTransform, CenterCrop 13 | from src.argus_models import SaltMetaModel 14 | from src.transforms import HorizontalFlip 15 | from src.utils import RLenc, make_dir 16 | from src import config 17 | 18 | EXPERIMENT_NAME = 'mos-fpn-lovasz-se-resnext50-001' 19 | FOLDS = list(range(config.N_FOLDS)) 20 | ORIG_IMAGE_SIZE = (101, 101) 21 | PRED_IMAGE_SIZE = (128, 128) 22 | TRANSFORM_MODE = 'crop' 23 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148_mos_emb_1.csv' 24 | IMAGES_NAME = '148' 25 | FOLDS_DIR = f'/workdir/data/experiments/{EXPERIMENT_NAME}' 26 | PREDICTION_DIR = f'/workdir/data/predictions/{EXPERIMENT_NAME}' 27 | make_dir(PREDICTION_DIR) 28 | 29 | SEGM_THRESH = 0.5 30 | PROB_THRESH = 0.5 31 | 32 | 33 | class Predictor: 34 | def __init__(self, model_path): 35 | self.model = load_model(model_path) 36 | self.model.nn_module.final = torch.nn.Sigmoid() 37 | self.model.nn_module.eval() 38 | 39 | self.depth_trns = SimpleDepthTransform() 40 | self.crop_trns = CenterCrop(ORIG_IMAGE_SIZE) 41 | self.trns = SaltTransform(PRED_IMAGE_SIZE, False, TRANSFORM_MODE) 42 | 43 | def __call__(self, image): 44 | tensor = self.depth_trns(image, 0) 45 | tensor = self.trns(tensor) 46 | tensor = tensor.unsqueeze(0).to(self.model.device) 47 | 48 | with torch.no_grad(): 49 | segm, prob = self.model.nn_module(tensor) 50 | segm = segm.cpu().numpy()[0][0] 51 | segm = self.crop_trns(segm) 52 | 53 | segm = (segm * 255).astype(np.uint8) 54 | prob = prob.item() 55 | 56 | return segm, prob 57 | 58 | 59 | def pred_val_fold(model_path, fold): 60 | predictor = Predictor(model_path) 61 | folds_df = pd.read_csv(TRAIN_FOLDS_PATH) 62 | fold_df = folds_df[folds_df.fold == fold] 63 | fold_prediction_dir = join(PREDICTION_DIR, f'fold_{fold}', 'val') 64 | make_dir(fold_prediction_dir) 65 | 66 | prob_dict = {'id': [], 'prob': []} 67 | for i, row in fold_df.iterrows(): 68 | image = cv2.imread(row.image_path, cv2.IMREAD_GRAYSCALE) 69 | segm, prob = predictor(image) 70 | segm_save_path = join(fold_prediction_dir, row.id + '.png') 71 | cv2.imwrite(segm_save_path, segm) 72 | 73 | prob_dict['id'].append(row.id) 74 | prob_dict['prob'].append(prob) 75 | 76 | prob_df = pd.DataFrame(prob_dict) 77 | prob_df.to_csv(join(fold_prediction_dir, 'probs.csv'), index=False) 78 | 79 | 80 | def pred_test_fold(model_path, fold): 81 | predictor = Predictor(model_path) 82 | prob_df = pd.read_csv(config.SAMPLE_SUBM_PATH) 83 | prob_df.rename(columns={'rle_mask': 'prob'}, inplace=True) 84 | 85 | fold_prediction_dir = join(PREDICTION_DIR, f'fold_{fold}', 'test') 86 | make_dir(fold_prediction_dir) 87 | 88 | for i, row in prob_df.iterrows(): 89 | image_path = join(config.TEST_DIR, 'images'+IMAGES_NAME, row.id + '.png') 90 | image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) 91 | segm, prob = predictor(image) 92 | row.prob = prob 93 | segm_save_path = join(fold_prediction_dir, row.id + '.png') 94 | cv2.imwrite(segm_save_path, segm) 95 | 96 | prob_df.to_csv(join(fold_prediction_dir, 'probs.csv'), index=False) 97 | 98 | 99 | def get_mean_probs_df(pred_dir): 100 | probs_df_lst = [] 101 | for i in range(len(FOLDS)): 102 | fold_dir = os.path.join(pred_dir, f'fold_{FOLDS[i]}') 103 | probs_df = pd.read_csv(join(fold_dir, 'test', 'probs.csv'), index_col='id') 104 | probs_df_lst.append(probs_df) 105 | 106 | mean_probs_df = probs_df_lst[0].copy() 107 | for probs_df in probs_df_lst[1:]: 108 | mean_probs_df.prob += probs_df.prob 109 | mean_probs_df.prob /= len(probs_df_lst) 110 | 111 | return mean_probs_df 112 | 113 | 114 | def make_mean_submission(): 115 | mean_probs_df = get_mean_probs_df(PREDICTION_DIR) 116 | 117 | sample_submition = pd.read_csv(config.SAMPLE_SUBM_PATH) 118 | 119 | for i, row in sample_submition.iterrows(): 120 | pred_name = row.id + '.png' 121 | segm_lst = [] 122 | for i in range(len(FOLDS)): 123 | fold_dir = os.path.join(PREDICTION_DIR, f'fold_{FOLDS[i]}') 124 | pred_path = join(fold_dir, 'test', pred_name) 125 | segm = cv2.imread(pred_path, cv2.IMREAD_GRAYSCALE) 126 | segm = segm.astype(np.float32) / 255 127 | segm_lst.append(segm) 128 | 129 | mean_segm = np.mean(segm_lst, axis=0) 130 | prob = mean_probs_df.loc[row.id].prob 131 | 132 | pred = mean_segm > SEGM_THRESH 133 | prob = int(prob > PROB_THRESH) 134 | pred = (pred * prob).astype(np.uint8) 135 | 136 | rle_mask = RLenc(pred) 137 | row.rle_mask = rle_mask 138 | 139 | sample_submition.to_csv(join(PREDICTION_DIR, f'{EXPERIMENT_NAME}-mean-subm.csv'), index=False) 140 | 141 | 142 | def get_best_model_path(dir_path): 143 | model_scores = [] 144 | for model_name in os.listdir(dir_path): 145 | score = re.search(r'-(\d+(?:\.\d+)?).pth', model_name) 146 | if score is not None: 147 | score = score.group(0)[1:-4] 148 | model_scores.append((model_name, score)) 149 | model_score = sorted(model_scores, key=lambda x: x[1]) 150 | best_model_name = model_score[-1][0] 151 | best_model_path = os.path.join(dir_path, best_model_name) 152 | return best_model_path 153 | 154 | 155 | if __name__ == "__main__": 156 | for i in range(len(FOLDS)): 157 | print("Predict fold", FOLDS[i]) 158 | fold_dir = os.path.join(FOLDS_DIR, f'fold_{FOLDS[i]}') 159 | best_model_path = get_best_model_path(fold_dir) 160 | print("Model path", best_model_path) 161 | print("Val predict") 162 | pred_val_fold(best_model_path, FOLDS[i]) 163 | print("Test predict") 164 | pred_test_fold(best_model_path, FOLDS[i]) 165 | 166 | print("Mean submission") 167 | make_mean_submission() 168 | -------------------------------------------------------------------------------- /pipeline/fpn-lovasz-se-resnext50-006-after-001/predict_folds.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join 3 | import re 4 | import cv2 5 | import numpy as np 6 | import pandas as pd 7 | 8 | from argus import load_model 9 | 10 | import torch 11 | 12 | from src.transforms import SimpleDepthTransform, SaltTransform, CenterCrop 13 | from src.argus_models import SaltMetaModel 14 | from src.transforms import HorizontalFlip 15 | from src.utils import RLenc, make_dir 16 | from src import config 17 | 18 | EXPERIMENT_NAME = 'fpn-lovasz-se-resnext50-006-after-001' 19 | N_FOLDS = 5 20 | FOLDS = list(range(N_FOLDS)) 21 | ORIG_IMAGE_SIZE = (101, 101) 22 | PRED_IMAGE_SIZE = (128, 128) 23 | TRANSFORM_MODE = 'crop' 24 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148.csv' 25 | IMAGES_NAME = '148' 26 | FOLDS_DIR = f'/workdir/data/experiments/{EXPERIMENT_NAME}' 27 | PREDICTION_DIR = f'/workdir/data/predictions/{EXPERIMENT_NAME}' 28 | make_dir(PREDICTION_DIR) 29 | 30 | SEGM_THRESH = 0.5 31 | PROB_THRESH = 0.5 32 | 33 | 34 | class Predictor: 35 | def __init__(self, model_path): 36 | self.model = load_model(model_path) 37 | self.model.nn_module.final = torch.nn.Sigmoid() 38 | self.model.nn_module.eval() 39 | 40 | self.depth_trns = SimpleDepthTransform() 41 | self.crop_trns = CenterCrop(ORIG_IMAGE_SIZE) 42 | self.trns = SaltTransform(PRED_IMAGE_SIZE, False, TRANSFORM_MODE) 43 | 44 | def __call__(self, image): 45 | tensor = self.depth_trns(image, 0) 46 | tensor = self.trns(tensor) 47 | tensor = tensor.unsqueeze(0).to(self.model.device) 48 | 49 | with torch.no_grad(): 50 | segm, prob = self.model.nn_module(tensor) 51 | segm = segm.cpu().numpy()[0][0] 52 | segm = self.crop_trns(segm) 53 | 54 | segm = (segm * 255).astype(np.uint8) 55 | prob = prob.item() 56 | 57 | return segm, prob 58 | 59 | 60 | def pred_val_fold(model_path, fold): 61 | predictor = Predictor(model_path) 62 | folds_df = pd.read_csv(TRAIN_FOLDS_PATH) 63 | fold_df = folds_df[folds_df.fold == fold] 64 | fold_prediction_dir = join(PREDICTION_DIR, f'fold_{fold}', 'val') 65 | make_dir(fold_prediction_dir) 66 | 67 | prob_dict = {'id': [], 'prob': []} 68 | for i, row in fold_df.iterrows(): 69 | image = cv2.imread(row.image_path, cv2.IMREAD_GRAYSCALE) 70 | segm, prob = predictor(image) 71 | segm_save_path = join(fold_prediction_dir, row.id + '.png') 72 | cv2.imwrite(segm_save_path, segm) 73 | 74 | prob_dict['id'].append(row.id) 75 | prob_dict['prob'].append(prob) 76 | 77 | prob_df = pd.DataFrame(prob_dict) 78 | prob_df.to_csv(join(fold_prediction_dir, 'probs.csv'), index=False) 79 | 80 | 81 | def pred_test_fold(model_path, fold): 82 | predictor = Predictor(model_path) 83 | prob_df = pd.read_csv(config.SAMPLE_SUBM_PATH) 84 | prob_df.rename(columns={'rle_mask': 'prob'}, inplace=True) 85 | 86 | fold_prediction_dir = join(PREDICTION_DIR, f'fold_{fold}', 'test') 87 | make_dir(fold_prediction_dir) 88 | 89 | for i, row in prob_df.iterrows(): 90 | image_path = join(config.TEST_DIR, 'images'+IMAGES_NAME, row.id + '.png') 91 | image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) 92 | segm, prob = predictor(image) 93 | row.prob = prob 94 | segm_save_path = join(fold_prediction_dir, row.id + '.png') 95 | cv2.imwrite(segm_save_path, segm) 96 | 97 | prob_df.to_csv(join(fold_prediction_dir, 'probs.csv'), index=False) 98 | 99 | 100 | def get_mean_probs_df(pred_dir): 101 | probs_df_lst = [] 102 | for i in range(len(FOLDS)): 103 | fold_dir = os.path.join(pred_dir, f'fold_{FOLDS[i]}') 104 | probs_df = pd.read_csv(join(fold_dir, 'test', 'probs.csv'), index_col='id') 105 | probs_df_lst.append(probs_df) 106 | 107 | mean_probs_df = probs_df_lst[0].copy() 108 | for probs_df in probs_df_lst[1:]: 109 | mean_probs_df.prob += probs_df.prob 110 | mean_probs_df.prob /= len(probs_df_lst) 111 | 112 | return mean_probs_df 113 | 114 | 115 | def make_mean_submission(): 116 | mean_probs_df = get_mean_probs_df(PREDICTION_DIR) 117 | 118 | sample_submition = pd.read_csv(config.SAMPLE_SUBM_PATH) 119 | 120 | for i, row in sample_submition.iterrows(): 121 | pred_name = row.id + '.png' 122 | segm_lst = [] 123 | for i in range(len(FOLDS)): 124 | fold_dir = os.path.join(PREDICTION_DIR, f'fold_{FOLDS[i]}') 125 | pred_path = join(fold_dir, 'test', pred_name) 126 | segm = cv2.imread(pred_path, cv2.IMREAD_GRAYSCALE) 127 | segm = segm.astype(np.float32) / 255 128 | segm_lst.append(segm) 129 | 130 | mean_segm = np.mean(segm_lst, axis=0) 131 | prob = mean_probs_df.loc[row.id].prob 132 | 133 | pred = mean_segm > SEGM_THRESH 134 | prob = int(prob > PROB_THRESH) 135 | pred = (pred * prob).astype(np.uint8) 136 | 137 | rle_mask = RLenc(pred) 138 | row.rle_mask = rle_mask 139 | 140 | sample_submition.to_csv(join(PREDICTION_DIR, f'{EXPERIMENT_NAME}-mean-subm.csv'), index=False) 141 | 142 | 143 | def get_best_model_path(dir_path): 144 | model_scores = [] 145 | for model_name in os.listdir(dir_path): 146 | score = re.search(r'-(\d+(?:\.\d+)?).pth', model_name) 147 | if score is not None: 148 | score = score.group(0)[1:-4] 149 | model_scores.append((model_name, score)) 150 | model_score = sorted(model_scores, key=lambda x: x[1]) 151 | best_model_name = model_score[-1][0] 152 | best_model_path = os.path.join(dir_path, best_model_name) 153 | return best_model_path 154 | 155 | 156 | if __name__ == "__main__": 157 | for i in range(len(FOLDS)): 158 | print("Predict fold", FOLDS[i]) 159 | fold_dir = os.path.join(FOLDS_DIR, f'fold_{FOLDS[i]}') 160 | best_model_path = get_best_model_path(fold_dir) 161 | print("Model path", best_model_path) 162 | print("Val predict") 163 | pred_val_fold(best_model_path, FOLDS[i]) 164 | print("Test predict") 165 | pred_test_fold(best_model_path, FOLDS[i]) 166 | 167 | print("Mean submission") 168 | make_mean_submission() 169 | -------------------------------------------------------------------------------- /notebooks/log_visualize.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import os\n", 10 | "\n", 11 | "from plotly import tools\n", 12 | "from plotly.offline import init_notebook_mode, plot, iplot\n", 13 | "import plotly.graph_objs as go\n", 14 | "\n", 15 | "\n", 16 | "init_notebook_mode(connected=True)" 17 | ] 18 | }, 19 | { 20 | "cell_type": "code", 21 | "execution_count": null, 22 | "metadata": {}, 23 | "outputs": [], 24 | "source": [ 25 | "def parce_log(log, key_names=['Train', 'Validation']):\n", 26 | " def fix_number(str_num):\n", 27 | " if str_num[-1] == ',':\n", 28 | " str_num = str_num[:-1]\n", 29 | " num_type = int\n", 30 | " if '.' in str_num:\n", 31 | " num_type = float\n", 32 | " return num_type(str_num)\n", 33 | " \n", 34 | " def fix_name(str_name):\n", 35 | " if str_name[-1] == ':':\n", 36 | " str_name = str_name[:-1]\n", 37 | " return str_name\n", 38 | "\n", 39 | " parsed_log = dict((key,[]) for key in key_names)\n", 40 | " for line in log:\n", 41 | " line = line.split()\n", 42 | " tine_type = line[3]\n", 43 | " if tine_type in key_names:\n", 44 | " keys = [fix_name(line[i]) for i in range(5, len(line), 2)]\n", 45 | " values = [fix_number(line[i]) for i in range(6, len(line), 2)]\n", 46 | " sample_dict = {k:v for (k,v) in zip(keys, values)}\n", 47 | " parsed_log[tine_type].append(sample_dict)\n", 48 | " return parsed_log\n", 49 | "\n", 50 | "def plot_data(log_list: list, log_id=None):\n", 51 | " x_name = 'Epoch'\n", 52 | " log_names = list(log_list[0].keys())\n", 53 | " data = [[] for _ in log_names]\n", 54 | "\n", 55 | " for sample in log_list:\n", 56 | " for i, name in enumerate(log_names):\n", 57 | " data[i].append(sample[name])\n", 58 | " \n", 59 | " plots = []\n", 60 | " for i in range(1, len(data)):\n", 61 | " if log_id is not None:\n", 62 | " name = log_names[i]+' '+str(log_id)\n", 63 | " else:\n", 64 | " name = log_names[i]\n", 65 | " plots.append(go.Scatter(x=data[0],\n", 66 | " y=data[i],\n", 67 | " mode='lines+markers',\n", 68 | " name=name))\n", 69 | " return plots\n", 70 | "\n", 71 | "def get_plot_data(log_path, key_names, log_id=None):\n", 72 | " data = []\n", 73 | " with open(log_path, 'r') as f: \n", 74 | " data = f.readlines()\n", 75 | "\n", 76 | " log = parce_log(data, key_names)\n", 77 | " plots = []\n", 78 | " for name in key_names:\n", 79 | " plots+= plot_data(log[name], log_id)\n", 80 | " return plots\n", 81 | "\n", 82 | "def draw_log(log_path, key_names, log_id=None):\n", 83 | " plots = get_plot_data(log_path, key_names, log_id)\n", 84 | " iplot(plots)\n", 85 | "\n", 86 | "def draw_folds(log_paths, key_names):\n", 87 | " paths = [log_path for log_path in log_paths]\n", 88 | " \n", 89 | " plots = [get_plot_data(log_path, key_names, i) for (i, log_path) in enumerate(paths)]\n", 90 | "\n", 91 | " for i, plot in enumerate(plots):\n", 92 | " layout = go.Layout(title='Fold '+str(i))\n", 93 | "\n", 94 | " fig = go.Figure(data=plot, layout=layout)\n", 95 | " iplot(fig)\n", 96 | "\n", 97 | "def compare_folds(log_paths, key_names):\n", 98 | " paths = [log_path for log_path in log_paths]\n", 99 | " plots = [get_plot_data(log_path, key_names, i) for (i, log_path) in enumerate(paths)]\n", 100 | "\n", 101 | " fig = tools.make_subplots(rows=len(plots[0]), cols=1, specs=[[{}], [{}], [{}]],\n", 102 | " shared_xaxes=True, shared_yaxes=False,\n", 103 | " vertical_spacing=0.1)\n", 104 | "\n", 105 | " for i, plot in enumerate(plots):\n", 106 | " for j in range(len(plot)):\n", 107 | " fig.append_trace(plot[j], j+1, 1)\n", 108 | " fig['layout'].update(height=300*len(plots), width=800, showlegend=False)\n", 109 | " iplot(fig)" 110 | ] 111 | }, 112 | { 113 | "cell_type": "code", 114 | "execution_count": null, 115 | "metadata": {}, 116 | "outputs": [], 117 | "source": [ 118 | "exp_path = '../data/experiments/densenet121-folds-001/'\n", 119 | "log_file = 'log.txt'\n", 120 | "key_names=['Train', 'Validation']\n", 121 | "\n", 122 | "folds_paths = [os.path.join(os.path.join(exp_path, fold_name), log_file) for fold_name in os.listdir(exp_path)\n", 123 | " if os.path.isdir(os.path.join(exp_path, fold_name))]" 124 | ] 125 | }, 126 | { 127 | "cell_type": "code", 128 | "execution_count": null, 129 | "metadata": {}, 130 | "outputs": [], 131 | "source": [ 132 | "draw_log(folds_paths[0], key_names)" 133 | ] 134 | }, 135 | { 136 | "cell_type": "code", 137 | "execution_count": null, 138 | "metadata": { 139 | "scrolled": false 140 | }, 141 | "outputs": [], 142 | "source": [ 143 | "draw_folds(folds_paths, key_names)" 144 | ] 145 | }, 146 | { 147 | "cell_type": "code", 148 | "execution_count": null, 149 | "metadata": {}, 150 | "outputs": [], 151 | "source": [ 152 | "compare_folds(folds_paths, key_names)" 153 | ] 154 | }, 155 | { 156 | "cell_type": "code", 157 | "execution_count": null, 158 | "metadata": {}, 159 | "outputs": [], 160 | "source": [] 161 | } 162 | ], 163 | "metadata": { 164 | "kernelspec": { 165 | "display_name": "Python 3", 166 | "language": "python", 167 | "name": "python3" 168 | }, 169 | "language_info": { 170 | "codemirror_mode": { 171 | "name": "ipython", 172 | "version": 3 173 | }, 174 | "file_extension": ".py", 175 | "mimetype": "text/x-python", 176 | "name": "python", 177 | "nbconvert_exporter": "python", 178 | "pygments_lexer": "ipython3", 179 | "version": "3.6.5" 180 | } 181 | }, 182 | "nbformat": 4, 183 | "nbformat_minor": 2 184 | } 185 | -------------------------------------------------------------------------------- /notebooks/train.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import os\n", 10 | "import cv2\n", 11 | "import numpy as np\n", 12 | "\n", 13 | "import argus\n", 14 | "from argus import Model\n", 15 | "from argus.callbacks import MonitorCheckpoint, EarlyStopping, LoggingToFile\n", 16 | "\n", 17 | "import torch\n", 18 | "import torch.nn as nn\n", 19 | "import torch.nn.functional as F\n", 20 | "from torchvision.transforms import ToTensor\n", 21 | "from torch.utils.data import DataLoader\n", 22 | "\n", 23 | "from src.dataset import SaltDataset\n", 24 | "from src.transforms import SimpleDepthTransform, DepthTransform, SaltTransform\n", 25 | "from src.argus_models import SaltProbModel, SaltMetaModel\n", 26 | "from src import config\n", 27 | "\n", 28 | "import matplotlib.pyplot as plt\n", 29 | "%matplotlib inline" 30 | ] 31 | }, 32 | { 33 | "cell_type": "code", 34 | "execution_count": null, 35 | "metadata": {}, 36 | "outputs": [], 37 | "source": [ 38 | "image_size = (128, 128)\n", 39 | "val_folds = [0]\n", 40 | "train_folds = [1, 2, 3, 4]\n", 41 | "train_batch_size = 64\n", 42 | "val_batch_size = 64" 43 | ] 44 | }, 45 | { 46 | "cell_type": "code", 47 | "execution_count": null, 48 | "metadata": {}, 49 | "outputs": [], 50 | "source": [ 51 | "depth_trns = SimpleDepthTransform()\n", 52 | "train_trns = SaltTransform(image_size, True, 'pad')\n", 53 | "val_trns = SaltTransform(image_size, False, 'pad')\n", 54 | "train_dataset = SaltDataset(config.TRAIN_FOLDS_PATH, train_folds, train_trns, depth_trns)\n", 55 | "val_dataset = SaltDataset(config.TRAIN_FOLDS_PATH, val_folds, val_trns, depth_trns)\n", 56 | "train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)\n", 57 | "val_loader = DataLoader(val_dataset, batch_size=val_batch_size, shuffle=False)" 58 | ] 59 | }, 60 | { 61 | "cell_type": "code", 62 | "execution_count": null, 63 | "metadata": {}, 64 | "outputs": [], 65 | "source": [ 66 | "# Draw a list of images in a row# Draw \n", 67 | "def draw(imgs):\n", 68 | " n = len(imgs) # Number of images in a row\n", 69 | " plt.figure(figsize=(7,n*7))\n", 70 | " for i in range(n):\n", 71 | " plt.subplot(1, n, i+1)\n", 72 | " plt.axis('off')\n", 73 | " plt.imshow(imgs[i])\n", 74 | " plt.show()" 75 | ] 76 | }, 77 | { 78 | "cell_type": "code", 79 | "execution_count": null, 80 | "metadata": {}, 81 | "outputs": [], 82 | "source": [ 83 | "n_images_to_draw = 3\n", 84 | "\n", 85 | "for img, trg in train_loader:\n", 86 | " for i in range(n_images_to_draw):\n", 87 | " img_i = img[i, 0, :, :].numpy()\n", 88 | " cumsum_i = img[i, 1, :, :].numpy()\n", 89 | " trg_i = trg[i, 0, :, :].numpy()\n", 90 | " draw([img_i, cumsum_i, trg_i])\n", 91 | " break" 92 | ] 93 | }, 94 | { 95 | "cell_type": "markdown", 96 | "metadata": {}, 97 | "source": [ 98 | "# Train prob" 99 | ] 100 | }, 101 | { 102 | "cell_type": "code", 103 | "execution_count": null, 104 | "metadata": {}, 105 | "outputs": [], 106 | "source": [ 107 | "params = {\n", 108 | " 'nn_module': ('DPNProbUnet', {\n", 109 | " 'num_classes': 1,\n", 110 | " 'num_channels': 3,\n", 111 | " 'encoder_name': 'dpn92',\n", 112 | " 'dropout': 0\n", 113 | " }),\n", 114 | " 'loss': ('FbBceProbLoss', {\n", 115 | " 'fb_weight': 0.95,\n", 116 | " 'fb_beta': 2,\n", 117 | " 'bce_weight': 0.9,\n", 118 | " 'prob_weight': 0.85\n", 119 | " }),\n", 120 | " 'prediction_transform': ('ProbOutputTransform', {\n", 121 | " 'segm_thresh': 0.5,\n", 122 | " 'prob_thresh': 0.5\n", 123 | " }),\n", 124 | " 'optimizer': ('Adam', {'lr': 0.0001}),\n", 125 | " 'device': 'cuda'\n", 126 | "}\n", 127 | "\n", 128 | "model = SaltMetaModel(params)" 129 | ] 130 | }, 131 | { 132 | "cell_type": "code", 133 | "execution_count": null, 134 | "metadata": { 135 | "scrolled": true 136 | }, 137 | "outputs": [], 138 | "source": [ 139 | "callbacks = [\n", 140 | " MonitorCheckpoint('/workdir/data/experiments/test_022', monitor='val_crop_iout', max_saves=3),\n", 141 | " EarlyStopping(monitor='val_crop_iout', patience=50),\n", 142 | " LoggingToFile('/workdir/data/experiments/test_022/log.txt')\n", 143 | "]\n", 144 | "\n", 145 | "model.fit(train_loader, \n", 146 | " val_loader=val_loader,\n", 147 | " max_epochs=1000,\n", 148 | " callbacks=callbacks,\n", 149 | " metrics=['crop_iout'])" 150 | ] 151 | }, 152 | { 153 | "cell_type": "code", 154 | "execution_count": null, 155 | "metadata": { 156 | "scrolled": true 157 | }, 158 | "outputs": [], 159 | "source": [ 160 | "from argus import load_model\n", 161 | "\n", 162 | "experiment_name = 'test_025'\n", 163 | "lr_steps = [\n", 164 | " (300, 0.0001),\n", 165 | " (300, 0.00003),\n", 166 | " (300, 0.00001),\n", 167 | " (300, 0.000003),\n", 168 | " (1000, 0.0000001)\n", 169 | "]\n", 170 | "\n", 171 | "for i, (epochs, lr) in enumerate(lr_steps):\n", 172 | " print(i, epochs, lr)\n", 173 | " if not i:\n", 174 | " model = SaltMetaModel(params)\n", 175 | " else:\n", 176 | " model = load_model(f'/workdir/data/experiments/{experiment_name}/model-last.pth')\n", 177 | " \n", 178 | " callbacks = [\n", 179 | " MonitorCheckpoint(f'/workdir/data/experiments/{experiment_name}', monitor='val_crop_iout', max_saves=2),\n", 180 | " EarlyStopping(monitor='val_crop_iout', patience=50),\n", 181 | " LoggingToFile(f'/workdir/data/experiments/{experiment_name}/log.txt')\n", 182 | " ] \n", 183 | " \n", 184 | " model.set_lr(lr)\n", 185 | " model.fit(train_loader, \n", 186 | " val_loader=val_loader,\n", 187 | " max_epochs=epochs,\n", 188 | " callbacks=callbacks,\n", 189 | " metrics=['crop_iout'])" 190 | ] 191 | } 192 | ], 193 | "metadata": { 194 | "kernelspec": { 195 | "display_name": "Python 3", 196 | "language": "python", 197 | "name": "python3" 198 | }, 199 | "language_info": { 200 | "codemirror_mode": { 201 | "name": "ipython", 202 | "version": 3 203 | }, 204 | "file_extension": ".py", 205 | "mimetype": "text/x-python", 206 | "name": "python", 207 | "nbconvert_exporter": "python", 208 | "pygments_lexer": "ipython3", 209 | "version": "3.6.5" 210 | } 211 | }, 212 | "nbformat": 4, 213 | "nbformat_minor": 2 214 | } 215 | -------------------------------------------------------------------------------- /predict_folds.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os.path import join 3 | import re 4 | import cv2 5 | import numpy as np 6 | import pandas as pd 7 | 8 | from argus import load_model 9 | 10 | import torch 11 | 12 | from src.transforms import SimpleDepthTransform, SaltTransform, CenterCrop 13 | from src.argus_models import SaltMetaModel 14 | from src.transforms import HorizontalFlip 15 | from src.utils import RLenc, make_dir 16 | from src import config 17 | 18 | EXPERIMENT_NAME = 'mos-fpn-lovasz-se-resnext50-001' 19 | FOLDS = list(range(config.N_FOLDS)) 20 | ORIG_IMAGE_SIZE = (101, 101) 21 | PRED_IMAGE_SIZE = (128, 128) 22 | TRANSFORM_MODE = 'crop' 23 | TRAIN_FOLDS_PATH = '/workdir/data/train_folds_148_mos_emb_1.csv' 24 | IMAGES_NAME = '148' 25 | FOLDS_DIR = f'/workdir/data/experiments/{EXPERIMENT_NAME}' 26 | PREDICTION_DIR = f'/workdir/data/predictions/{EXPERIMENT_NAME}' 27 | make_dir(PREDICTION_DIR) 28 | 29 | SEGM_THRESH = 0.5 30 | PROB_THRESH = 0.5 31 | 32 | 33 | class Predictor: 34 | def __init__(self, model_path): 35 | self.model = load_model(model_path) 36 | self.model.nn_module.final = torch.nn.Sigmoid() # 37 | self.model.nn_module.eval() 38 | 39 | self.depth_trns = SimpleDepthTransform() 40 | self.crop_trns = CenterCrop(ORIG_IMAGE_SIZE) 41 | self.trns = SaltTransform(PRED_IMAGE_SIZE, False, TRANSFORM_MODE) 42 | 43 | def __call__(self, image): 44 | tensor = self.depth_trns(image, 0) 45 | tensor = self.trns(tensor) 46 | tensor = tensor.unsqueeze(0).to(self.model.device) 47 | 48 | with torch.no_grad(): 49 | segm, prob = self.model.nn_module(tensor) 50 | segm = segm.cpu().numpy()[0][0] 51 | segm = self.crop_trns(segm) 52 | 53 | segm = (segm * 255).astype(np.uint8) 54 | prob = prob.item() 55 | 56 | return segm, prob 57 | 58 | 59 | class TtaPredictor: 60 | def __init__(self, model_path): 61 | self.model = load_model(model_path) 62 | self.model.nn_module.eval() 63 | 64 | self.depth_trns = SimpleDepthTransform() 65 | self.crop_trns = CenterCrop(ORIG_IMAGE_SIZE) 66 | self.trns = SaltTransform(PRED_IMAGE_SIZE, False, TRANSFORM_MODE) 67 | 68 | self.flip = HorizontalFlip() 69 | 70 | def forward_tta(self, image): 71 | image = self.depth_trns(image) 72 | images = [ 73 | image, 74 | self.flip(image) 75 | ] 76 | images = [self.trns(img) for img in images] 77 | tensor = torch.stack(images, dim=0) 78 | return tensor 79 | 80 | def backward_tta(self, output): 81 | segms, probs = output 82 | segms = segms.cpu().numpy() 83 | probs = probs.cpu().numpy() 84 | mean_prob = float(np.mean(probs)) 85 | 86 | segm_lst = [] 87 | for i in range(segms.shape[0]): 88 | segm = segms[i, 0] 89 | segm = self.crop_trns(segm) 90 | if i == 1: 91 | segm = self.flip(segm) 92 | segm_lst.append(segm) 93 | 94 | mean_segm = np.mean(segm_lst, axis=0) 95 | mean_segm = (mean_segm * 255).astype(np.uint8) 96 | return mean_segm, mean_prob 97 | 98 | def __call__(self, image): 99 | tensor = self.forward_tta(image) 100 | tensor = tensor.to(self.model.device) 101 | 102 | with torch.no_grad(): 103 | output = self.model.nn_module(tensor) 104 | segm, prob = self.backward_tta(output) 105 | return segm, prob 106 | 107 | 108 | def pred_val_fold(model_path, fold): 109 | predictor = Predictor(model_path) 110 | folds_df = pd.read_csv(TRAIN_FOLDS_PATH) 111 | fold_df = folds_df[folds_df.fold == fold] 112 | fold_prediction_dir = join(PREDICTION_DIR, f'fold_{fold}', 'val') 113 | make_dir(fold_prediction_dir) 114 | 115 | prob_dict = {'id': [], 'prob': []} 116 | for i, row in fold_df.iterrows(): 117 | image = cv2.imread(row.image_path, cv2.IMREAD_GRAYSCALE) 118 | segm, prob = predictor(image) 119 | segm_save_path = join(fold_prediction_dir, row.id + '.png') 120 | cv2.imwrite(segm_save_path, segm) 121 | 122 | prob_dict['id'].append(row.id) 123 | prob_dict['prob'].append(prob) 124 | 125 | prob_df = pd.DataFrame(prob_dict) 126 | prob_df.to_csv(join(fold_prediction_dir, 'probs.csv'), index=False) 127 | 128 | 129 | def pred_test_fold(model_path, fold): 130 | predictor = Predictor(model_path) 131 | prob_df = pd.read_csv(config.SAMPLE_SUBM_PATH) 132 | prob_df.rename(columns={'rle_mask': 'prob'}, inplace=True) 133 | 134 | fold_prediction_dir = join(PREDICTION_DIR, f'fold_{fold}', 'test') 135 | make_dir(fold_prediction_dir) 136 | 137 | for i, row in prob_df.iterrows(): 138 | image_path = join(config.TEST_DIR, 'images'+IMAGES_NAME, row.id + '.png') 139 | image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) 140 | segm, prob = predictor(image) 141 | row.prob = prob 142 | segm_save_path = join(fold_prediction_dir, row.id + '.png') 143 | cv2.imwrite(segm_save_path, segm) 144 | 145 | prob_df.to_csv(join(fold_prediction_dir, 'probs.csv'), index=False) 146 | 147 | 148 | def get_mean_probs_df(pred_dir): 149 | probs_df_lst = [] 150 | for i in range(len(FOLDS)): 151 | fold_dir = os.path.join(pred_dir, f'fold_{FOLDS[i]}') 152 | probs_df = pd.read_csv(join(fold_dir, 'test', 'probs.csv'), index_col='id') 153 | probs_df_lst.append(probs_df) 154 | 155 | mean_probs_df = probs_df_lst[0].copy() 156 | for probs_df in probs_df_lst[1:]: 157 | mean_probs_df.prob += probs_df.prob 158 | mean_probs_df.prob /= len(probs_df_lst) 159 | 160 | return mean_probs_df 161 | 162 | 163 | def make_mean_submission(): 164 | mean_probs_df = get_mean_probs_df(PREDICTION_DIR) 165 | 166 | sample_submition = pd.read_csv(config.SAMPLE_SUBM_PATH) 167 | 168 | for i, row in sample_submition.iterrows(): 169 | pred_name = row.id + '.png' 170 | segm_lst = [] 171 | for i in range(len(FOLDS)): 172 | fold_dir = os.path.join(PREDICTION_DIR, f'fold_{FOLDS[i]}') 173 | pred_path = join(fold_dir, 'test', pred_name) 174 | segm = cv2.imread(pred_path, cv2.IMREAD_GRAYSCALE) 175 | segm = segm.astype(np.float32) / 255 176 | segm_lst.append(segm) 177 | 178 | mean_segm = np.mean(segm_lst, axis=0) 179 | prob = mean_probs_df.loc[row.id].prob 180 | 181 | pred = mean_segm > SEGM_THRESH 182 | prob = int(prob > PROB_THRESH) 183 | pred = (pred * prob).astype(np.uint8) 184 | 185 | rle_mask = RLenc(pred) 186 | row.rle_mask = rle_mask 187 | 188 | sample_submition.to_csv(join(PREDICTION_DIR, f'{EXPERIMENT_NAME}-mean-subm.csv'), index=False) 189 | 190 | 191 | def get_best_model_path(dir_path): 192 | model_scores = [] 193 | for model_name in os.listdir(dir_path): 194 | score = re.search(r'-(\d+(?:\.\d+)?).pth', model_name) 195 | if score is not None: 196 | score = score.group(0)[1:-4] 197 | model_scores.append((model_name, score)) 198 | model_score = sorted(model_scores, key=lambda x: x[1]) 199 | best_model_name = model_score[-1][0] 200 | best_model_path = os.path.join(dir_path, best_model_name) 201 | return best_model_path 202 | 203 | 204 | if __name__ == "__main__": 205 | for i in range(len(FOLDS)): 206 | print("Predict fold", FOLDS[i]) 207 | fold_dir = os.path.join(FOLDS_DIR, f'fold_{FOLDS[i]}') 208 | best_model_path = get_best_model_path(fold_dir) 209 | print("Model path", best_model_path) 210 | print("Val predict") 211 | pred_val_fold(best_model_path, FOLDS[i]) 212 | print("Test predict") 213 | pred_test_fold(best_model_path, FOLDS[i]) 214 | 215 | print("Mean submission") 216 | make_mean_submission() 217 | -------------------------------------------------------------------------------- /src/lovasz.py: -------------------------------------------------------------------------------- 1 | """ 2 | Lovasz-Softmax and Jaccard hinge loss in PyTorch 3 | Maxim Berman 2018 ESAT-PSI KU Leuven (MIT License) 4 | """ 5 | 6 | from __future__ import print_function, division 7 | 8 | import torch 9 | from torch.autograd import Variable 10 | import torch.nn.functional as F 11 | import numpy as np 12 | try: 13 | from itertools import ifilterfalse 14 | except ImportError: # py3k 15 | from itertools import filterfalse 16 | 17 | 18 | def lovasz_grad(gt_sorted): 19 | """ 20 | Computes gradient of the Lovasz extension w.r.t sorted errors 21 | See Alg. 1 in paper 22 | """ 23 | p = len(gt_sorted) 24 | gts = gt_sorted.sum() 25 | intersection = gts - gt_sorted.float().cumsum(0) 26 | union = gts + (1 - gt_sorted).float().cumsum(0) 27 | jaccard = 1. - intersection / union 28 | if p > 1: # cover 1-pixel case 29 | jaccard[1:p] = jaccard[1:p] - jaccard[0:-1] 30 | return jaccard 31 | 32 | 33 | def iou_binary(preds, labels, EMPTY=1., ignore=None, per_image=True): 34 | """ 35 | IoU for foreground class 36 | binary: 1 foreground, 0 background 37 | """ 38 | if not per_image: 39 | preds, labels = (preds,), (labels,) 40 | ious = [] 41 | for pred, label in zip(preds, labels): 42 | intersection = ((label == 1) & (pred == 1)).sum() 43 | union = ((label == 1) | ((pred == 1) & (label != ignore))).sum() 44 | if not union: 45 | iou = EMPTY 46 | else: 47 | iou = float(intersection) / union 48 | ious.append(iou) 49 | iou = mean(ious) # mean accross images if per_image 50 | return 100 * iou 51 | 52 | 53 | def iou(preds, labels, C, EMPTY=1., ignore=None, per_image=False): 54 | """ 55 | Array of IoU for each (non ignored) class 56 | """ 57 | if not per_image: 58 | preds, labels = (preds,), (labels,) 59 | ious = [] 60 | for pred, label in zip(preds, labels): 61 | iou = [] 62 | for i in range(C): 63 | if i != ignore: # The ignored label is sometimes among predicted classes (ENet - CityScapes) 64 | intersection = ((label == i) & (pred == i)).sum() 65 | union = ((label == i) | ((pred == i) & (label != ignore))).sum() 66 | if not union: 67 | iou.append(EMPTY) 68 | else: 69 | iou.append(float(intersection) / union) 70 | ious.append(iou) 71 | ious = map(mean, zip(*ious)) # mean accross images if per_image 72 | return 100 * np.array(ious) 73 | 74 | 75 | # --------------------------- BINARY LOSSES --------------------------- 76 | 77 | 78 | def lovasz_hinge(logits, labels, per_image=True, ignore=None): 79 | """ 80 | Binary Lovasz hinge loss 81 | logits: [B, H, W] Variable, logits at each pixel (between -\infty and +\infty) 82 | labels: [B, H, W] Tensor, binary ground truth masks (0 or 1) 83 | per_image: compute the loss per image instead of per batch 84 | ignore: void class id 85 | """ 86 | if per_image: 87 | loss = mean(lovasz_hinge_flat(*flatten_binary_scores(log.unsqueeze(0), lab.unsqueeze(0), ignore)) 88 | for log, lab in zip(logits, labels)) 89 | else: 90 | loss = lovasz_hinge_flat(*flatten_binary_scores(logits, labels, ignore)) 91 | return loss 92 | 93 | 94 | def lovasz_hinge_flat(logits, labels): 95 | """ 96 | Binary Lovasz hinge loss 97 | logits: [P] Variable, logits at each prediction (between -\infty and +\infty) 98 | labels: [P] Tensor, binary ground truth labels (0 or 1) 99 | ignore: label to ignore 100 | """ 101 | if len(labels) == 0: 102 | # only void pixels, the gradients should be 0 103 | return logits.sum() * 0. 104 | signs = 2. * labels.float() - 1. 105 | errors = (1. - logits * Variable(signs)) 106 | errors_sorted, perm = torch.sort(errors, dim=0, descending=True) 107 | perm = perm.data 108 | gt_sorted = labels[perm] 109 | grad = lovasz_grad(gt_sorted) 110 | loss = torch.dot(F.elu(errors_sorted) + 1, Variable(grad)) 111 | return loss 112 | 113 | 114 | def flatten_binary_scores(scores, labels, ignore=None): 115 | """ 116 | Flattens predictions in the batch (binary case) 117 | Remove labels equal to 'ignore' 118 | """ 119 | scores = scores.view(-1) 120 | labels = labels.view(-1) 121 | if ignore is None: 122 | return scores, labels 123 | valid = (labels != ignore) 124 | vscores = scores[valid] 125 | vlabels = labels[valid] 126 | return vscores, vlabels 127 | 128 | 129 | class StableBCELoss(torch.nn.modules.Module): 130 | def __init__(self): 131 | super(StableBCELoss, self).__init__() 132 | def forward(self, input, target): 133 | neg_abs = - input.abs() 134 | loss = input.clamp(min=0) - input * target + (1 + neg_abs.exp()).log() 135 | return loss.mean() 136 | 137 | 138 | def binary_xloss(logits, labels, ignore=None): 139 | """ 140 | Binary Cross entropy loss 141 | logits: [B, H, W] Variable, logits at each pixel (between -\infty and +\infty) 142 | labels: [B, H, W] Tensor, binary ground truth masks (0 or 1) 143 | ignore: void class id 144 | """ 145 | logits, labels = flatten_binary_scores(logits, labels, ignore) 146 | loss = StableBCELoss()(logits, Variable(labels.float())) 147 | return loss 148 | 149 | 150 | # --------------------------- MULTICLASS LOSSES --------------------------- 151 | 152 | 153 | def lovasz_softmax(probas, labels, only_present=False, per_image=False, ignore=None): 154 | """ 155 | Multi-class Lovasz-Softmax loss 156 | probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1) 157 | labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1) 158 | only_present: average only on classes present in ground truth 159 | per_image: compute the loss per image instead of per batch 160 | ignore: void class labels 161 | """ 162 | if per_image: 163 | loss = mean(lovasz_softmax_flat(*flatten_probas(prob.unsqueeze(0), lab.unsqueeze(0), ignore), only_present=only_present) 164 | for prob, lab in zip(probas, labels)) 165 | else: 166 | loss = lovasz_softmax_flat(*flatten_probas(probas, labels, ignore), only_present=only_present) 167 | return loss 168 | 169 | 170 | def lovasz_softmax_flat(probas, labels, only_present=False): 171 | """ 172 | Multi-class Lovasz-Softmax loss 173 | probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1) 174 | labels: [P] Tensor, ground truth labels (between 0 and C - 1) 175 | only_present: average only on classes present in ground truth 176 | """ 177 | C = probas.size(1) 178 | losses = [] 179 | for c in range(C): 180 | fg = (labels == c).float() # foreground for class c 181 | if only_present and fg.sum() == 0: 182 | continue 183 | errors = (Variable(fg) - probas[:, c]).abs() 184 | errors_sorted, perm = torch.sort(errors, 0, descending=True) 185 | perm = perm.data 186 | fg_sorted = fg[perm] 187 | losses.append(torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted)))) 188 | return mean(losses) 189 | 190 | 191 | def flatten_probas(probas, labels, ignore=None): 192 | """ 193 | Flattens predictions in the batch 194 | """ 195 | B, C, H, W = probas.size() 196 | probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C) # B * H * W, C = P, C 197 | labels = labels.view(-1) 198 | if ignore is None: 199 | return probas, labels 200 | valid = (labels != ignore) 201 | vprobas = probas[valid.nonzero().squeeze()] 202 | vlabels = labels[valid] 203 | return vprobas, vlabels 204 | 205 | def xloss(logits, labels, ignore=None): 206 | """ 207 | Cross entropy loss 208 | """ 209 | return F.cross_entropy(logits, Variable(labels), ignore_index=255) 210 | 211 | 212 | # --------------------------- HELPER FUNCTIONS --------------------------- 213 | 214 | def mean(l, ignore_nan=False, empty=0): 215 | """ 216 | nanmean compatible with generators. 217 | """ 218 | l = iter(l) 219 | if ignore_nan: 220 | l = ifilterfalse(np.isnan, l) 221 | try: 222 | n = 1 223 | acc = next(l) 224 | except StopIteration: 225 | if empty == 'raise': 226 | raise ValueError('Empty mean') 227 | return empty 228 | for n, v in enumerate(l, 2): 229 | acc += v 230 | if n == 1: 231 | return acc 232 | return acc / n 233 | -------------------------------------------------------------------------------- /src/nick_zoo/unet_flex.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from src.nick_zoo.resnet_blocks import resnet34 5 | from src.nick_zoo.layers import DecoderBlock, BasicBlock, Bottleneck 6 | from src.nick_zoo.utils import class_for_name 7 | 8 | nonlinearity = nn.ReLU 9 | 10 | 11 | class UNetFlexProb(nn.Module): 12 | block_dict = { 13 | 'basic': BasicBlock, 14 | 'bottleneck': Bottleneck 15 | } 16 | 17 | def __init__(self, num_classes, num_channels=3, blocks=resnet34, 18 | final='softmax', dropout_2d=0.5, deflation=4, 19 | is_deconv=True, first_last_conv=[64, 32], 20 | skip_dropout=False, xavier=False, pretrain=None, 21 | pretrain_layers=None, use_first_pool=True, 22 | return_enc=False): 23 | super(UNetFlexProb, self).__init__() 24 | assert num_channels > 0, "Incorrect num channels" 25 | assert final in ['softmax', 'sigmoid'],\ 26 | "Incorrect output type. Should be 'softmax' or 'sigmoid'" 27 | assert len(first_last_conv)==2, "'first_last_conv' - list with number of the first and the last convs" 28 | self.inplanes = first_last_conv[0] 29 | self.dropout_2d = dropout_2d 30 | self.use_first_pool = use_first_pool 31 | self.return_enc = return_enc 32 | 33 | # Initial convolutions 34 | self.firstconv = nn.Conv2d(num_channels, self.inplanes, kernel_size=(7, 7), 35 | stride=(2, 2), padding=(3, 3), bias=False) 36 | self.firstbn = nn.BatchNorm2d(self.inplanes) 37 | self.firstrelu = nonlinearity(inplace=True) 38 | self.firstmaxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 39 | 40 | # Encoders 41 | self.encoders = nn.ModuleList([]) 42 | encoders_n = [] 43 | for block in blocks: 44 | self.encoders.append(self._make_encoder(self.block_dict[block['type']], 45 | block['n_blocks'], 46 | block['n_filters'], 47 | block['stride'])) 48 | encoders_n.append(self.inplanes) 49 | 50 | # Decoders 51 | self.decoders = nn.ModuleList([DecoderBlock(2*encoders_n[0], encoders_n[0], 52 | is_deconv, deflation=deflation)]) 53 | for i in range(1, len(encoders_n)-1): 54 | self.decoders.append(DecoderBlock(2*encoders_n[i], encoders_n[i-1], 55 | is_deconv, deflation=deflation)) 56 | self.decoders.append(DecoderBlock(encoders_n[-1], encoders_n[-2], 57 | is_deconv, deflation=deflation)) 58 | print("Actual number of filters:", encoders_n) 59 | self.encoders_n = encoders_n 60 | if skip_dropout and dropout_2d is not None: 61 | self.skip_dropouts = nn.ModuleList([nn.Dropout2d(p=dropout_2d) 62 | for _ in range(len(self.decoders)-1)]) 63 | else: 64 | self.skip_dropouts = None 65 | 66 | # Final classifier 67 | if self.use_first_pool: 68 | if is_deconv: 69 | self.finalupscale = nn.ConvTranspose2d(encoders_n[0], 70 | first_last_conv[1], 3, 71 | stride=2) 72 | self.finalconv2 = nn.Conv2d(first_last_conv[1], first_last_conv[1], 3) 73 | else: 74 | self.finalupscale = nn.Upsample(scale_factor=2, mode='bilinear', 75 | align_corners=False) 76 | self.finalconv2 = nn.Conv2d(encoders_n[0], first_last_conv[1], 3) 77 | else: 78 | self.finalconv2 = nn.Conv2d(encoders_n[0], first_last_conv[1], 3, padding=1) 79 | self.finalrelu1 = nonlinearity(inplace=True) 80 | self.finaldropout = nn.Dropout2d(p=dropout_2d) if dropout_2d\ 81 | is not None else None 82 | self.finalrelu2 = nonlinearity(inplace=True) 83 | if self.use_first_pool: 84 | self.finalconv3 = nn.Conv2d(first_last_conv[1], num_classes, 2, padding=1) 85 | else: 86 | self.finalconv3 = nn.Conv2d(first_last_conv[1], num_classes, 3, padding=1) 87 | # Prob branch 88 | self.pool = nn.MaxPool2d(2, 2) 89 | self.prob_avg_pool = nn.AdaptiveAvgPool2d(1) 90 | self.prob_fc = nn.Linear(encoders_n[-1], 1) 91 | self.sigmoid = nn.Sigmoid() 92 | 93 | if final=='softmax': 94 | self.final = nn.Softmax(dim=1) 95 | else: 96 | self.final = nn.Sigmoid() 97 | 98 | for m in self.modules(): 99 | if isinstance(m, nn.Conv2d): 100 | if xavier: 101 | nn.init.xavier_uniform_(m.weight) 102 | else: 103 | # Default resnet initialization from torchvision 104 | nn.init.kaiming_normal_(m.weight, mode='fan_out', 105 | nonlinearity='relu') 106 | elif isinstance(m, nn.BatchNorm2d): 107 | nn.init.constant_(m.weight, 1) 108 | nn.init.constant_(m.bias, 0) 109 | 110 | if pretrain is not None and pretrain_layers is not None: 111 | use_pretrain_resnet(self, pretrain, pretrain_layers) 112 | 113 | 114 | def _make_encoder(self, block, n_blocks, planes, stride=1): 115 | downsample = None 116 | if stride != 1 or self.inplanes != planes * block.expansion: 117 | downsample = nn.Sequential( 118 | nn.Conv2d(self.inplanes, planes * block.expansion, 119 | kernel_size=1, stride=stride, bias=False), 120 | nn.BatchNorm2d(planes * block.expansion), 121 | ) 122 | 123 | layers = [] 124 | layers.append(block(self.inplanes, planes, stride, downsample)) 125 | self.inplanes = planes * block.expansion 126 | for i in range(1, n_blocks): 127 | layers.append(block(self.inplanes, planes)) 128 | 129 | return nn.Sequential(*layers) 130 | 131 | def forward(self, x): 132 | # Encoder 133 | x = self.firstconv(x) 134 | x = self.firstbn(x) 135 | x = self.firstrelu(x) 136 | if self.use_first_pool: 137 | x = self.firstmaxpool(x) 138 | 139 | enc = [] 140 | for encoder in self.encoders: 141 | x = encoder(x) 142 | enc.append(x) 143 | 144 | # Decoder with Skip Connections 145 | for i in range(len(self.decoders)-1, 0, -1): 146 | x = self.decoders[i](x) 147 | if self.skip_dropouts is not None: 148 | x = torch.cat([x, self.skip_dropouts[i-1](enc[i-1])], dim=1) 149 | else: 150 | x = torch.cat([x, enc[i-1]], dim=1) 151 | 152 | x = self.decoders[0](x) 153 | 154 | # Final Classification 155 | if self.use_first_pool: 156 | x = self.finalupscale(x) 157 | x = self.finalrelu1(x) 158 | if self.finaldropout is not None: 159 | x = self.finaldropout(x) 160 | x = self.finalconv2(x) 161 | x = self.finalrelu2(x) 162 | x = self.finalconv3(x) 163 | x = self.final(x) 164 | 165 | prob = self.prob_avg_pool(enc[-1]) 166 | prob = prob.view(prob.size(0), -1) 167 | prob = self.prob_fc(prob) 168 | prob = self.sigmoid(prob) 169 | prob = prob.view(-1) 170 | if self.return_enc: 171 | return x, prob, enc 172 | else: 173 | return x, prob 174 | 175 | 176 | def use_pretrain_resnet(model, encoder, layers): 177 | # Unsafe function, use wisely! 178 | assert encoder in ['resnet18', 'resnet34', 'resnet50',\ 179 | 'resnet101', 'resnet152'],\ 180 | "Incorrect encoder type" 181 | resnet = class_for_name("torchvision.models", encoder)\ 182 | (pretrained=True) 183 | resnet_layers = [resnet.layer1, resnet.layer2, 184 | resnet.layer3, resnet.layer4] 185 | if layers[0]: 186 | model.firstconv = resnet.conv1 187 | model.firstbn = resnet.bn1 188 | model.firstrelu = resnet.relu 189 | model.firstmaxpool = resnet.maxpool 190 | 191 | if len(layers) > 1: 192 | for i, layer in enumerate(layers[1:]): 193 | if layer: 194 | model.encoders[i] = resnet_layers[i] 195 | -------------------------------------------------------------------------------- /mosaic/create_mosaic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import time 4 | import tqdm 5 | import json 6 | import numpy as np 7 | from collections import defaultdict 8 | from multiprocessing import Pool 9 | 10 | img_train_dir = '../../data/train/images/' 11 | img_test_dir = '../../data/test/images/' 12 | images_names = ['14b50d2c26', '31f29ada31', '26c403ce6d', '3eef3e190f', 'b2d4c41f68', '14c2834524', 'cb36e2e2ae', '971ff26c0d', 'a0891f7e3e', 13 | '0b9874fd4f', '476bde9dbf', 'ad068f08d1', '4d33311a1e', '7dfdf6eeb8', '27838f7f46', 'f27a3ca60e', 'da324cbdde', 'd0bdfc3217', '393048bc6c', 14 | 'eb84dfdf18', 'c100d07aa9', '3d0f8c81ce', 'a667ba911e', '973a14e68f', 'cf5be83765', 'b174e802d6', '3768877c25', '1d6f51de05', 'c24e12d0cf', 15 | 'f40bfdeb45', '55f1ea73b3', 'f91af813d6', '8d0c065461', 'f96f608c5a', '4b366ad75d', 'de15d35ebc', 'e5c26c8634', 16 | 'f8ba90c991', '59dabe134f', '5f723327ef', 'd9d7c77e13', 'b0beff63b4', '4aa5b85fb1', '4b6bc93e6c', 'cb6c6c15ce', '595deae139', '5420263191', 17 | '1986fcadda', '0c14b0fe28', 'd639f4c3ae', '499237cf6c', 'dc71cdae70', '9472731dfd', '6230e29f51', '34f944baa4', '3ab326340b', 18 | '7cf980df36', '2fecaf1c54', 'c2ec2c9de4', '8a58fe70f8', 'd3e4630469', '9a91290888', 'c5d6528a35', '3be0c0be8e', 'fca4d42846', '70b7e1c459', 19 | 'abd9ba8120', 'cbc5bd4021', '52ae147d60', '8a6090c9ec', '7205efa791', '70afc514d2', '4ca0fd980f', '668cf61ce0', 'f448297676', 'd96f1003cf', 20 | 'c4f2855630', '44d28db2dc', '94500b7464', 'ea39579316', '753ac9f3ed', 'a0e6d8b0a7', '7db7a5fb8f', '8299b97f38', '73d3d08aff', '095473ab35', 21 | '61b6c452ad', '6d69267940', 'f9f6588f79', '54629440c9', 'cd90c05864', '0bbaa6d56a', '5299cd2c33', 'f069247787', 22 | ] 23 | THRES = 5 24 | diff_order = 1 25 | # Outputs 26 | mosaic_dict_path = '../../data/mosaic/mosaic1.json' 27 | clusters_dict_path = '../../data/mosaic/clusters1.json' 28 | img_mosaics_path = '../../data/mosaic/imgs2/' 29 | mean_hist_path = '../../data/mosaic/mean_hist.npy' 30 | save_mtrx = True 31 | left_mtrx_path = '../../data/mosaic/left1.npy' 32 | top_mtrx_path = '../../data/mosaic/top1.npy' 33 | 34 | images_names = os.listdir(img_train_dir) 35 | #images_names += os.listdir(img_test_dir) 36 | 37 | mean_hist = np.load(mean_hist_path) 38 | mean_hist[0] = mean_hist[-1] 39 | 40 | def make_descriptor(img): 41 | edges = [img[:, 0], img[:, -1], img[0, :], img[-1, :]] 42 | return [(np.diff(edge.astype(np.int), n=diff_order), edge) 43 | for edge in edges] 44 | 45 | def metric(desc1, desc2): 46 | return np.mean((desc1[0]-desc2[0])**2) + np.mean((desc1[1]-desc2[1])**2) 47 | 48 | def create_mosaic_dict(images_paths, tiles, tile_to_cluster): 49 | tiles_dicts = {} 50 | for i in range(n_samples): 51 | tiles_dicts[i] = { 52 | 'path': images_paths[i], 53 | 'cluster': tile_to_cluster[tiles[i][0]], 54 | 'neighbours': tiles[i][1:] 55 | } 56 | return tiles_dicts 57 | 58 | def hist_match(source, hist): 59 | """ 60 | Adjust the pixel values of a grayscale image such that its histogram 61 | matches a given one 62 | 63 | Arguments: 64 | ----------- 65 | source: np.ndarray 66 | Image to transform; the histogram is computed over the flattened 67 | array 68 | template: np.ndarray 69 | Template histogram 70 | Returns: 71 | ----------- 72 | matched: np.ndarray 73 | The transformed output image 74 | """ 75 | 76 | oldshape = source.shape 77 | source = source.ravel() 78 | 79 | # get the set of unique pixel values and their corresponding indices and 80 | # counts 81 | s_values, bin_idx, s_counts = np.unique(source, return_inverse=True, 82 | return_counts=True) 83 | t_values = np.linspace(0, 255, 256).astype(np.int) 84 | 85 | # take the cumsum of the counts and normalize by the number of pixels to 86 | # get the empirical cumulative distribution functions for the source and 87 | # template images (maps pixel value --> quantile) 88 | s_quantiles = np.cumsum(s_counts).astype(np.float64) 89 | s_quantiles /= s_quantiles[-1] 90 | t_quantiles = np.cumsum(hist).astype(np.float64) 91 | t_quantiles /= t_quantiles[-1] 92 | 93 | # interpolate linearly to find the pixel values in the template image 94 | # that correspond most closely to the quantiles in the source image 95 | interp_t_values = np.interp(s_quantiles, t_quantiles, t_values) 96 | 97 | return interp_t_values[bin_idx].reshape(oldshape) 98 | 99 | n_samples = len(images_names) 100 | 101 | images_paths = [] 102 | for img_name in images_names: 103 | if img_name[-3:] != 'png': 104 | train_path = os.path.join(img_train_dir, img_name+'.png') 105 | test_path = os.path.join(img_test_dir, img_name+'.png') 106 | else: 107 | train_path = os.path.join(img_train_dir, img_name) 108 | test_path = os.path.join(img_test_dir, img_name) 109 | if os.path.exists(train_path): 110 | images_paths.append(train_path) 111 | elif os.path.exists(test_path): 112 | images_paths.append(test_path) 113 | 114 | # Check if everything is ok 115 | print("Samples:", n_samples, "Ok:", len(images_paths)==len(images_names)) 116 | 117 | imgs = [hist_match(cv2.imread(img_path, cv2.IMREAD_GRAYSCALE), mean_hist) 118 | for img_path in images_paths] 119 | #imgs = [cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) 120 | #for img_path in images_paths] 121 | 122 | 123 | 124 | c_time = time.time() 125 | descriptors = [make_descriptor(img) for img in imgs] 126 | 127 | def calc_metrics(ij): 128 | i, j = ij 129 | if i != j: 130 | left_metr = metric(descriptors[i][0], descriptors[j][1]) 131 | top_metr = metric(descriptors[i][2], descriptors[j][3]) 132 | else: 133 | left_metr = top_metr = 10000 # Just a big number 134 | return left_metr, top_metr 135 | 136 | left_matrix = np.zeros((n_samples, n_samples)) 137 | top_matrix = np.zeros((n_samples, n_samples)) 138 | print("Descriptors matrix creation") 139 | with Pool(processes=12) as pool: 140 | for i in tqdm.tqdm(range(len(descriptors))): 141 | metr = pool.imap(calc_metrics, [(i, j) for j in 142 | range(len(descriptors))]) 143 | for j, m in enumerate(metr): 144 | left_matrix[i, j] = m[0] 145 | top_matrix[i, j] = m[1] 146 | 147 | if save_mtrx: 148 | np.save(left_mtrx_path, left_matrix) 149 | np.save(top_mtrx_path, top_matrix) 150 | 151 | tiles = [[i, None, None, None, None] for i in range(n_samples)] 152 | 153 | print("Create left neighbours") 154 | left_thres = np.median(left_matrix, axis=1)/THRES 155 | left_idx = np.argsort(left_matrix, axis=1)[:,0] 156 | 157 | for i in range(n_samples): 158 | j = int(left_idx[i]) 159 | if left_matrix[i, j] < left_thres[i]: 160 | tiles[i][1] = j 161 | tiles[j][3] = i 162 | 163 | print("Create top neighbours") 164 | top_thres = np.median(top_matrix, axis=1)/THRES 165 | top_idx = np.argsort(top_matrix, axis=1)[:,0] 166 | 167 | for i in range(n_samples): 168 | j = int(top_idx[i]) 169 | if top_matrix[i, j] < top_thres[i]: 170 | tiles[i][2] = j 171 | tiles[j][4] = i 172 | 173 | 174 | clusters = [] 175 | n_clusters = 1 176 | tile_to_cluster = defaultdict(int) 177 | print("Clusters assigning") 178 | for tile in tqdm.tqdm(tiles): 179 | tile_cluster = tile_to_cluster[tile[0]] 180 | if tile_cluster == 0: 181 | tile_to_cluster[tile[0]] = n_clusters 182 | n_clusters += 1 183 | tile_cluster = tile_to_cluster[tile[0]] 184 | for idx in tile[1:]: 185 | if idx is not None: 186 | if tile_to_cluster[idx] == 0: 187 | tile_to_cluster[idx] = tile_cluster 188 | else: 189 | # TODO Think about collisions solution 190 | pass 191 | 192 | print(time.time()-c_time, "sec passed") 193 | 194 | tiles_dicts = create_mosaic_dict(images_paths, tiles, tile_to_cluster) 195 | with open(mosaic_dict_path, 'w') as fout: 196 | json.dump(tiles_dicts, fout, indent=2) 197 | 198 | 199 | img_size = (101,101) 200 | cls_id = 3 201 | 202 | def coords_for_open_list(i, j, neighbours, open_list, close_list): 203 | to_open_list = [] 204 | for n, k in enumerate(neighbours): 205 | if k is not None and k not in close_list: 206 | close_list.append(k) 207 | i_k = i 208 | j_k = j 209 | if n == 0: 210 | i_k -= 1 # left 211 | elif n == 1: 212 | j_k -= 1 # top 213 | elif n == 2: 214 | i_k += 1 # right 215 | else: 216 | j_k += 1 # bottom 217 | to_open_list.append((i_k, j_k, k)) 218 | return open_list+to_open_list, close_list 219 | 220 | print("Assemble mosaics") 221 | def create_mosaics(tiles_dicts): 222 | 223 | def coords_ij(ij, lim): 224 | i, j = ij 225 | i_max, j_max, i_min, j_min = lim 226 | x_min = (i-i_min)*img_size[1] 227 | x_max = (i+1-i_min)*img_size[1] 228 | y_min = (j-j_min)*img_size[0] 229 | y_max = (j+1-j_min)*img_size[0] 230 | return x_min, x_max, y_min, y_max 231 | 232 | def add_img(mosaic, ij, lim, el): 233 | i, j = ij 234 | i_max, j_max, i_min, j_min = lim 235 | tile = tiles_dicts[el[2]] 236 | if el[1] < j_min: 237 | mosaic = np.pad(mosaic, ((img_size[0], 0), (0, 0)), mode='constant') 238 | j_min -= 1 239 | if el[1] > j_max: 240 | mosaic = np.pad(mosaic, ((0, img_size[0]), (0, 0)), mode='constant') 241 | j_max += 1 242 | if el[0] < i_min: 243 | mosaic = np.pad(mosaic, ((0, 0), (img_size[0], 0)), mode='constant') 244 | i_min -= 1 245 | if el[0] > i_max: 246 | mosaic = np.pad(mosaic, ((0, 0), (0, img_size[0])), mode='constant') 247 | i_max += 1 248 | #print('res mosaic', mosaic.shape) 249 | x_min, x_max, y_min, y_max = coords_ij(el[:2], 250 | (i_max, j_max, i_min, j_min)) 251 | #print("Add", el[:2], el[2], [x_min, x_max, y_min, y_max]) 252 | mosaic[y_min:y_max, x_min:x_max] = hist_match(cv2.imread(tile['path'], 253 | cv2.IMREAD_GRAYSCALE), mean_hist) 254 | return mosaic, i_max, j_max, i_min, j_min 255 | 256 | 257 | mosaics = [] 258 | dicts_to_save = [] 259 | while len(tiles_dicts) > 0: 260 | min_idx = min(tiles_dicts.keys()) 261 | tile = tiles_dicts[min_idx] 262 | del tiles_dicts[min_idx] 263 | mosaics.append(cv2.imread(tile['path'], 264 | cv2.IMREAD_GRAYSCALE)) 265 | i = 1 266 | j = 1 267 | i_max = 1 268 | j_max = 1 269 | i_min = 1 270 | j_min = 1 271 | open_list = [] 272 | close_list = [min_idx] 273 | open_list, close_list = coords_for_open_list(i, j, tile['neighbours'], 274 | open_list, close_list) 275 | dicts_to_save.append([]) 276 | dicts_to_save[-1].append({'i': i, 'j': j, 'path': tile['path']}) 277 | while len(open_list) > 0: 278 | el = open_list.pop(0) 279 | if el[2] in tiles_dicts.keys(): 280 | mosaics[-1], i_max, j_max, i_min, j_min = add_img(mosaics[-1], (i, j), (i_max, j_max, i_min, j_min), el) 281 | i,j = el[:2] 282 | open_list, close_list = coords_for_open_list(i, j, tiles_dicts[el[2]]['neighbours'], 283 | open_list, close_list) 284 | dicts_to_save[-1].append({'i': i, 'j': j, 'path': tiles_dicts[el[2]]['path']}) 285 | del tiles_dicts[el[2]] 286 | #print(open_list, close_list) 287 | return mosaics, dicts_to_save 288 | 289 | mosaics, dicts_to_save = create_mosaics(tiles_dicts) 290 | 291 | for i, mosaic in enumerate(mosaics): 292 | cv2.imwrite(os.path.join(img_mosaics_path, str(i)+'.png'), 293 | mosaic) 294 | 295 | with open(clusters_dict_path, 'w') as fout: 296 | json.dump(dicts_to_save, fout, indent=2) 297 | -------------------------------------------------------------------------------- /mosaic/mosaic_api.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pandas as pd 3 | import cv2 4 | import os 5 | from os.path import join 6 | from collections import defaultdict 7 | 8 | from src.utils import filename_without_ext 9 | from src.config import TRAIN_DIR, TEST_DIR 10 | 11 | IMAGE_SHAPE = (101, 101) 12 | 13 | FONT = cv2.FONT_HERSHEY_SIMPLEX 14 | COL_WHITE = (255,255,255) 15 | 16 | def stack_images(images_array): 17 | img_rows = [] 18 | for row in images_array: 19 | img_row = [] 20 | for img in row: 21 | img_row.append(img) 22 | img_rows.append(np.concatenate(img_row, axis=1)) 23 | return np.concatenate(img_rows, axis=0) 24 | 25 | 26 | def split_mosaic(mosaic): 27 | img_array = [] 28 | for row_img in np.split(mosaic, mosaic.shape[0]//101, axis=0): 29 | img_row_array = [] 30 | for img in np.split(row_img, mosaic.shape[1]//101, axis=1): 31 | img_row_array.append(img) 32 | img_array.append(img_row_array) 33 | return img_array 34 | 35 | 36 | class Mosaic: 37 | def __init__(self, raw_mosaic): 38 | 39 | i_min = min([p['i'] for p in raw_mosaic]) 40 | j_min = min([p['j'] for p in raw_mosaic]) 41 | 42 | i_max = -np.inf 43 | j_max = -np.inf 44 | 45 | mosaic = [] 46 | for raw_piece in raw_mosaic: 47 | piece = dict() 48 | piece['i'] = raw_piece['i'] - i_min 49 | piece['j'] = raw_piece['j'] - j_min 50 | piece['id'] = raw_piece['id'] 51 | mosaic.append(piece) 52 | 53 | if piece['i'] > i_max: 54 | i_max = piece['i'] 55 | if piece['j'] > j_max: 56 | j_max = piece['j'] 57 | 58 | id2index = dict() 59 | mosaic_array = np.full((j_max + 1, i_max + 1), fill_value=None, dtype=object) 60 | for piece in mosaic: 61 | mosaic_array[piece['j'], piece['i']] = piece['id'] 62 | id2index[piece['id']] = piece['j'], piece['i'] 63 | 64 | self.pieces = mosaic 65 | self.array = mosaic_array 66 | self.id2index = id2index 67 | self.ids = set(self.id2index.keys()) 68 | 69 | def get_neighbor(self, id, i_shift, j_shift): 70 | assert id in self.id2index, f'Mosaic have not piece with id {id}' 71 | 72 | j, i = self.id2index[id] 73 | neig_j = j + j_shift 74 | neig_i = i + i_shift 75 | if neig_j < 0 or neig_j >= self.array.shape[0]: 76 | return None 77 | if neig_i < 0 or neig_i >= self.array.shape[1]: 78 | return None 79 | return self.array[neig_j, neig_i] 80 | 81 | def get_neighbors(self, id): 82 | assert id in self.id2index, f'Mosaic have not piece with id {id}' 83 | 84 | return np.array(np.matrix([ 85 | [self.get_neighbor(id, -1, -1), self.get_neighbor(id, 0, -1), self.get_neighbor(id, 1, -1)], 86 | [self.get_neighbor(id, -1, 0), id, self.get_neighbor(id, 1, 0)], 87 | [self.get_neighbor(id, -1, 1), self.get_neighbor(id, 0, 1), self.get_neighbor(id, 1, 1)] 88 | ])) 89 | 90 | def __contains__(self, item): 91 | return item in self.id2index 92 | 93 | def __repr__(self): 94 | string = "" 95 | for j in range(self.array.shape[0]): 96 | for i in range(self.array.shape[1]): 97 | string += f'{self.array[j, i]} ' 98 | string += '\n' 99 | return string 100 | 101 | def __len__(self): 102 | return len(self.pieces) 103 | 104 | 105 | class Mosaics: 106 | def __init__(self, mosaic_csv_path): 107 | self.mosaics_csv_path = mosaic_csv_path 108 | 109 | mosaic_df = pd.read_csv(mosaic_csv_path) 110 | raw_mosaic_dict = defaultdict(list) 111 | 112 | for i, row in mosaic_df.iterrows(): 113 | raw_mosaic_dict[row.mosaic_id].append({ 114 | 'i': row.x, 115 | 'j': row.y, 116 | 'id': row.id 117 | }) 118 | 119 | self.mosaic_id2mosaic = dict() 120 | self.id2mosaic_id = dict() 121 | self.id2mosaic = dict() 122 | self.ids = set() 123 | for mosaic_id, raw_mosaic in raw_mosaic_dict.items(): 124 | mosaic = Mosaic(raw_mosaic) 125 | for id in mosaic.ids: 126 | self.id2mosaic_id[id] = mosaic_id 127 | self.id2mosaic[id] = mosaic 128 | self.mosaic_id2mosaic[mosaic_id] = mosaic 129 | self.ids.update(mosaic.ids) 130 | 131 | def __contains__(self, id): 132 | return id in self.ids 133 | 134 | def __getitem__(self, mosaic_id): 135 | return self.mosaic_id2mosaic[mosaic_id] 136 | 137 | 138 | class SaltData: 139 | def __init__(self, 140 | train_dir=TRAIN_DIR, 141 | test_dir=TEST_DIR, 142 | pred_dir=None, 143 | images_dir_name='images', 144 | masks_dir_name='masks', 145 | mosaic_csv_path=None): 146 | test_images_dir = join(test_dir, images_dir_name) 147 | test_images_names = os.listdir(test_images_dir) 148 | self.pred_dir = pred_dir 149 | train_images_dir = join(train_dir, images_dir_name) 150 | train_masks_dir = join(train_dir, masks_dir_name) 151 | train_images_names = os.listdir(train_images_dir) 152 | self.empty_image = np.zeros(IMAGE_SHAPE, dtype=np.uint8) 153 | self.full_image = np.full(IMAGE_SHAPE, fill_value=255, dtype=np.uint8) 154 | 155 | img_paths = [join(test_images_dir, name) for name in test_images_names] 156 | img_paths += [join(train_images_dir, name) for name in train_images_names] 157 | mask_paths = [join(train_masks_dir, name) for name in train_images_names] 158 | self.id2pred = {} 159 | if self.pred_dir is not None: 160 | pred_images_names = os.listdir(self.pred_dir) 161 | pred_mask_paths = [join(self.pred_dir, name) 162 | for name in pred_images_names] 163 | self.id2pred = {filename_without_ext(pred_mask_path): 164 | cv2.imread(pred_mask_path, cv2.IMREAD_GRAYSCALE) 165 | for pred_mask_path in pred_mask_paths} 166 | self.id2image = {filename_without_ext(img_path): cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) 167 | for img_path in img_paths} 168 | self.id2mask = {filename_without_ext(mask_path): cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) 169 | for mask_path in mask_paths} 170 | self.id2pred_cor = {} 171 | self.train_ids = set([filename_without_ext(n) for n in train_images_names]) 172 | self.test_ids = set([filename_without_ext(n) for n in test_images_names]) 173 | 174 | # Remove empty images from train 175 | for id in list(self.train_ids): 176 | image = self.id2image[id] 177 | if not np.sum(image): 178 | del self.id2image[id] 179 | self.train_ids.remove(id) 180 | 181 | self.ids = self.train_ids | self.test_ids 182 | 183 | self.mosaics = mosaic_csv_path 184 | if mosaic_csv_path is not None: 185 | self.mosaics = Mosaics(mosaic_csv_path) 186 | 187 | def in_train(self, id): 188 | assert id in self.ids 189 | return id in self.train_ids 190 | 191 | def in_mosaics(self, id): 192 | assert id in self.ids 193 | return id in self.mosaics 194 | 195 | def get_neighbors(self, id): 196 | if id in self.mosaics: 197 | mosaic = self.mosaics.id2mosaic[id] 198 | return mosaic.get_neighbors(id) 199 | else: 200 | return np.array(np.matrix([ 201 | [None, None, None], 202 | [None, id, None], 203 | [None, None, None], 204 | ])) 205 | 206 | def add_pred_cor(self, id, pred_cor): 207 | self.id2pred_cor[id] = pred_cor 208 | 209 | def draw_visualization(self, id_array, draw_names=False): 210 | img_array = [] 211 | for row in id_array: 212 | img_row = [] 213 | for id in row: 214 | if id is None: 215 | img = np.stack([self.empty_image]*3, axis=2) 216 | else: 217 | img = self.id2image[id].astype(int) 218 | img = np.stack([img] * 3, axis=2) 219 | if id in self.id2mask: 220 | img[:, :, 1] += (self.id2mask[id] * 0.25).astype(int) 221 | img[:, :, 2] += 25 222 | if id in self.id2pred: 223 | img[:, :, 0] += (self.id2pred[id] * 0.25).astype(int) 224 | if id in self.id2pred_cor: 225 | img[:, :, 0] += (self.id2pred_cor[id] * 0.25).astype(int) 226 | img[:, :, 1] += (self.id2pred_cor[id] * 0.25).astype(int) 227 | img = np.clip(img, 0, 255) 228 | img = img.astype(np.uint8) 229 | if draw_names: 230 | cv2.putText(img, str(id), (5, 60), 231 | FONT, 0.45, COL_WHITE, 1) 232 | img_row.append(img.astype(np.uint8)) 233 | img_array.append(img_row) 234 | return stack_images(img_array) 235 | 236 | def get_stacked_images(self, id_array, return_unknown_mask=False): 237 | img_array = [] 238 | unknown_mask_array = [] 239 | for row in id_array: 240 | img_row = [] 241 | unknown_mask_row = [] 242 | for id in row: 243 | if id is None: 244 | img = self.empty_image 245 | unknown_mask = self.full_image 246 | else: 247 | img = self.id2image[id] 248 | unknown_mask = self.empty_image 249 | img_row.append(img) 250 | unknown_mask_row.append(unknown_mask) 251 | img_array.append(img_row) 252 | unknown_mask_array.append(unknown_mask_row) 253 | if return_unknown_mask: 254 | return stack_images(img_array), stack_images(unknown_mask_array) 255 | return stack_images(img_array) 256 | 257 | def get_stacked_masks(self, id_array, return_unknown_mask=False): 258 | mask_array = [] 259 | unknown_mask_array = [] 260 | for row in id_array: 261 | mask_row = [] 262 | unknown_mask_row = [] 263 | for id in row: 264 | if id is None: 265 | mask = self.empty_image 266 | unknown_mask = self.full_image 267 | elif id in self.id2mask: 268 | # In train 269 | mask = self.id2mask[id] 270 | unknown_mask = self.empty_image 271 | else: 272 | # In test 273 | mask = self.empty_image 274 | unknown_mask = self.full_image 275 | mask_row.append(mask) 276 | unknown_mask_row.append(unknown_mask) 277 | mask_array.append(mask_row) 278 | unknown_mask_array.append(unknown_mask_row) 279 | if return_unknown_mask: 280 | return stack_images(mask_array), stack_images(unknown_mask_array) 281 | return stack_images(mask_array) 282 | 283 | def get_pred_stacked_masks(self, id_array, return_unknown_mask=False): 284 | mask_array = [] 285 | unknown_mask_array = [] 286 | for row in id_array: 287 | mask_row = [] 288 | unknown_mask_row = [] 289 | for id in row: 290 | if id is None: 291 | mask = self.empty_image 292 | unknown_mask = self.full_image 293 | elif id in self.id2mask: 294 | # In train 295 | mask = self.id2mask[id] 296 | unknown_mask = self.empty_image 297 | else: 298 | # In test 299 | mask = self.id2pred[id] 300 | unknown_mask = self.full_image 301 | mask_row.append(mask) 302 | unknown_mask_row.append(unknown_mask) 303 | mask_array.append(mask_row) 304 | unknown_mask_array.append(unknown_mask_row) 305 | if return_unknown_mask: 306 | return stack_images(mask_array), stack_images(unknown_mask_array) 307 | return stack_images(mask_array) 308 | 309 | def get_pred_masks(self, id_array, return_unknown_mask=False): 310 | mask_array = [] 311 | unknown_mask_array = [] 312 | for row in id_array: 313 | mask_row = [] 314 | unknown_mask_row = [] 315 | for id in row: 316 | if id is None: 317 | mask = self.empty_image 318 | unknown_mask = self.full_image 319 | elif id in self.id2mask: 320 | # In train 321 | mask = self.id2mask[id] 322 | unknown_mask = self.empty_image 323 | else: 324 | # In test 325 | mask = self.id2pred[id] 326 | unknown_mask = self.full_image 327 | mask_row.append(mask) 328 | unknown_mask_row.append(unknown_mask) 329 | mask_array.append(mask_row) 330 | unknown_mask_array.append(unknown_mask_row) 331 | if return_unknown_mask: 332 | return mask_array, unknown_mask_array 333 | return mask_array 334 | -------------------------------------------------------------------------------- /src/transforms.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import random 3 | import numpy as np 4 | import pandas as pd 5 | 6 | import cv2 7 | 8 | cv2.setNumThreads(0) 9 | 10 | 11 | def img_size(image: np.ndarray): 12 | return (image.shape[1], image.shape[0]) 13 | 14 | 15 | def image_cumsum(img): 16 | cumsum_img = np.float32(img).cumsum(axis=0) 17 | cumsum_img -= cumsum_img.min() 18 | if cumsum_img.max() > 0: 19 | cumsum_img /= cumsum_img.max() 20 | cumsum_img *= 255 21 | cumsum_img = np.clip(cumsum_img, 0, 255) 22 | cumsum_img = cumsum_img.astype(np.uint8) 23 | return cumsum_img 24 | 25 | 26 | def get_samples_for_aug(folds_path): 27 | images_lst = [] 28 | train_folds_df = pd.read_csv(folds_path) 29 | 30 | for i, row in train_folds_df.iterrows(): 31 | mask = cv2.imread(row.mask_path, cv2.IMREAD_GRAYSCALE) 32 | if np.max(mask) == 0: 33 | image = cv2.imread(row.image_path, cv2.IMREAD_GRAYSCALE) 34 | images_lst.append(image) 35 | print("Images for background augmentation:", len(images_lst)) 36 | return images_lst 37 | 38 | 39 | def img_crop(img, box): 40 | return img[box[1]:box[3], box[0]:box[2]] 41 | 42 | 43 | def random_crop(img, size): 44 | tw = size[0] 45 | th = size[1] 46 | w, h = img_size(img) 47 | if ((w - tw) > 0) and ((h - th) > 0): 48 | x1 = random.randint(0, w - tw) 49 | y1 = random.randint(0, h - th) 50 | else: 51 | x1 = 0 52 | y1 = 0 53 | img_return = img_crop(img, (x1, y1, x1 + tw, y1 + th)) 54 | return img_return, x1, y1 55 | 56 | 57 | class ProbOutputTransform: 58 | def __init__(self, segm_thresh=0.5, prob_thresh=0.5, size=None): 59 | self.segm_thresh = segm_thresh 60 | self.prob_thresh = prob_thresh 61 | self.crop = None 62 | if size is not None: 63 | self.crop = CenterCrop(size) 64 | 65 | def __call__(self, preds): 66 | segms, probs = preds 67 | preds = segms > self.segm_thresh 68 | probs = probs > self.prob_thresh 69 | preds = preds * probs.view(-1, 1, 1, 1) 70 | if self.crop is not None: 71 | preds = self.crop(preds) 72 | return preds 73 | 74 | 75 | class OneOf: 76 | def __init__(self, transforms, prob=0.5): 77 | self.transforms = transforms 78 | self.prob = prob 79 | 80 | def __call__(self, image, mask=None): 81 | if mask is None: 82 | if random.random() < self.prob: 83 | transform = random.choice(self.transforms) 84 | image = transform(image) 85 | return image 86 | else: 87 | if random.random() < self.prob: 88 | transform = random.choice(self.transforms) 89 | image, mask = transform(image, mask) 90 | return image, mask 91 | 92 | 93 | class UseWithProb: 94 | def __init__(self, transform, prob=.5): 95 | self.transform = transform 96 | self.prob = prob 97 | 98 | def __call__(self, image, mask=None): 99 | if mask is None: 100 | if random.random() < self.prob: 101 | image = self.transform(image) 102 | return image 103 | else: 104 | if random.random() < self.prob: 105 | image, mask = self.transform(image, mask) 106 | return image, mask 107 | 108 | 109 | class Compose(object): 110 | def __init__(self, transforms): 111 | self.transforms = transforms 112 | 113 | def __call__(self, image, mask=None): 114 | if mask is None: 115 | for trns in self.transforms: 116 | image = trns(image) 117 | return image 118 | else: 119 | for trns in self.transforms: 120 | image, mask = trns(image, mask) 121 | return image, mask 122 | 123 | 124 | class GaussNoise: 125 | def __init__(self, sigma_sq): 126 | self.sigma_sq = sigma_sq 127 | 128 | def __call__(self, img, mask=None): 129 | if self.sigma_sq > 0.0: 130 | img = self._gauss_noise(img, 131 | np.random.uniform(0, self.sigma_sq)) 132 | if mask is None: 133 | return img 134 | return img, mask 135 | 136 | def _gauss_noise(self, img, sigma_sq): 137 | img = img.astype(np.uint32) 138 | w, h, c = img.shape 139 | gauss = np.random.normal(0, sigma_sq, (h, w)) 140 | gauss = gauss.reshape(h, w) 141 | img = img + np.stack([gauss for i in range(c)], axis=2) 142 | img = np.clip(img, 0, 255) 143 | img = img.astype(np.uint8) 144 | return img 145 | 146 | 147 | class SpeckleNoise: 148 | def __init__(self, sigma_sq): 149 | self.sigma_sq = sigma_sq 150 | 151 | def __call__(self, img, mask=None): 152 | if self.sigma_sq > 0.0: 153 | img = self._speckle_noise(img, 154 | np.random.uniform(0, self.sigma_sq)) 155 | if mask is None: 156 | return img 157 | return img, mask 158 | 159 | def _speckle_noise(self, img, sigma_sq): 160 | sigma_sq /= 255 161 | img = img.astype(np.uint32) 162 | w, h, c = img.shape 163 | gauss = np.random.normal(0, sigma_sq, (h, w)) 164 | gauss = gauss.reshape(h, w) 165 | img = img + np.stack([gauss for i in range(c)], axis=2) * img 166 | img = np.clip(img, 0, 255) 167 | img = img.astype(np.uint8) 168 | return img 169 | 170 | 171 | class RandomCrop: 172 | def __init__(self, size): 173 | self.size = size 174 | 175 | def __call__(self, img, mask): 176 | img, x1, y1 = random_crop(img, self.size) 177 | if mask is None: 178 | return img 179 | mask = img_crop(mask, (x1, y1, x1 + self.size[0], 180 | y1 + self.size[1])) 181 | return img, mask 182 | 183 | 184 | class Flip: 185 | def __init__(self, flip_code): 186 | assert flip_code in [0, 1] 187 | self.flip_code = flip_code 188 | 189 | def __call__(self, img, mask=None): 190 | img = cv2.flip(img, self.flip_code) 191 | if mask is None: 192 | return img 193 | mask = cv2.flip(mask, self.flip_code) 194 | return img, mask 195 | 196 | 197 | class HorizontalFlip(Flip): 198 | def __init__(self): 199 | super().__init__(1) 200 | 201 | 202 | class RandomBg: 203 | def __init__(self, input_path, max_w): 204 | # TODO Update if not use the SimpleDepthTransform 205 | self.bg_imgs = get_samples_for_aug(input_path) 206 | assert max_w < 1.0, "Weight of the augmentet background should be < 1.0" 207 | self.max_w = max_w 208 | 209 | def __call__(self, img, mask=None): 210 | w = np.random.uniform(0, self.max_w) 211 | n = np.random.randint(0, len(self.bg_imgs)) 212 | bg = self.bg_imgs[n] 213 | img = cv2.addWeighted(img, 1.0 - w, np.stack([bg, bg, bg], axis=2), w, 0) 214 | if mask is None: 215 | return img 216 | return img, mask 217 | 218 | 219 | class PadToSize: 220 | def __init__(self, size, random_center=False, mode='reflect'): 221 | self.size = np.array(size, int) 222 | self.random_center = random_center 223 | self.mode = mode 224 | 225 | def __call__(self, img, mask=None): 226 | diff = self.size - img.shape[:2] 227 | assert np.all(diff >= 0) 228 | 229 | if self.random_center: 230 | pad_left = [] 231 | for i in range(len(diff)): 232 | if diff[i]: 233 | pad_left.append(np.random.randint(0, diff[i])) 234 | else: 235 | pad_left.append(0) 236 | else: 237 | pad_left = diff // 2 238 | pad_right = diff - pad_left 239 | pad_width = ((pad_left[0], pad_right[0]), (pad_left[1], pad_right[1])) 240 | if self.mode == 'constant': 241 | pad_img = np.pad(img, pad_width + ((0, 0),), self.mode, constant_values=0) 242 | else: 243 | pad_img = np.pad(img, pad_width + ((0, 0),), self.mode) 244 | if mask is None: 245 | return pad_img 246 | pad_mask = np.pad(mask, pad_width, self.mode) 247 | return pad_img, pad_mask 248 | 249 | 250 | class CenterCrop(object): 251 | def __init__(self, size): 252 | self.size = size 253 | 254 | def __call__(self, img, mask=None): 255 | w, h = img.shape[1], img.shape[0] 256 | tw, th = self.size 257 | if w == tw and h == th: 258 | if mask is None: 259 | return img 260 | else: 261 | return img, mask 262 | 263 | x1 = (w - tw) // 2 264 | y1 = (h - th) // 2 265 | 266 | crop_img = img_crop(img, (x1, y1, x1 + tw, y1 + th)) 267 | 268 | if mask is None: 269 | return crop_img 270 | crop_mask = img_crop(mask, (x1, y1, x1 + tw, y1 + th)) 271 | return crop_img, crop_mask 272 | 273 | 274 | class RandomHorizontalScale: 275 | def __init__(self, max_scale, interpolation=cv2.INTER_NEAREST): 276 | assert max_scale >= 1.0, "RandomHorizontalScale works as upscale only" 277 | self.max_scale = max_scale 278 | self.interpolation = interpolation 279 | 280 | def __call__(self, image, mask=None): 281 | if self.max_scale > 1.0: 282 | scale_factor = np.random.uniform(1.0, self.max_scale) 283 | resized_image = self._scale_img(image, scale_factor) 284 | if mask is None: 285 | return resize_image 286 | resized_mask = self._scale_img(mask, scale_factor) 287 | return resized_image, resized_mask 288 | else: 289 | if mask is None: 290 | return image 291 | return image, mask 292 | 293 | def _scale_img(self, img, factor): 294 | h, w = img.shape[:2] 295 | new_w = round(w * factor) 296 | img = cv2.resize(img, dsize=(new_w, h), 297 | interpolation=self.interpolation) 298 | if new_w > w: 299 | random_x = np.random.randint(0, new_w - w) 300 | else: 301 | random_x = 0 302 | return img[:, random_x:(random_x + w)] 303 | 304 | 305 | class Blur: 306 | def __init__(self, ksize=(5, 1), sigma_x=20): 307 | self.ksize = ksize 308 | self.sigma_x = sigma_x 309 | 310 | def __call__(self, image, mask=None): 311 | blured_image = cv2.GaussianBlur(image, self.ksize, self.sigma_x) 312 | if mask is None: 313 | return blured_image 314 | return blured_image, mask 315 | 316 | 317 | class Scale: 318 | def __init__(self, size, interpolation=cv2.INTER_LINEAR): 319 | self.size = tuple(size) 320 | self.interpolation = interpolation 321 | 322 | def __call__(self, image, mask=None): 323 | resize_image = cv2.resize(image, self.size, interpolation=self.interpolation) 324 | if mask is None: 325 | return resize_image 326 | resize_mask = cv2.resize(mask, self.size, interpolation=cv2.INTER_NEAREST) 327 | return resize_image, resize_mask 328 | 329 | 330 | class SimpleDepthTransform: 331 | def __call__(self, image, depth=None): 332 | input = np.stack([image, image, image], axis=2) 333 | return input 334 | 335 | 336 | class DepthTransform: 337 | def __call__(self, image, depth): 338 | cumsum_img = image_cumsum(image) 339 | depth_img = np.full_like(image, depth // 4) 340 | input = np.stack([image, cumsum_img, depth_img], axis=2) 341 | return input 342 | 343 | 344 | class ImageToTensor: 345 | def __init__(self, coord_channels=False): 346 | self.coord_channels = coord_channels 347 | 348 | def add_coord_channels(self, image_tensor): 349 | _, h, w = image_tensor.size() 350 | for row, const in enumerate(np.linspace(0, 1, h)): 351 | image_tensor[1, row, :] = const 352 | for col, const in enumerate(np.linspace(0, 1, w)): 353 | image_tensor[2, :, col] = const 354 | return image_tensor 355 | 356 | def __call__(self, image): 357 | image = np.moveaxis(image, -1, 0) 358 | image = image.astype(np.float32) / 255.0 359 | image = torch.from_numpy(image) 360 | if self.coord_channels: 361 | image = self.add_coord_channels(image) 362 | return image 363 | 364 | 365 | class MaskToTensor: 366 | def __call__(self, mask): 367 | mask[mask > 0] = 1 368 | return mask.astype(np.float32)[np.newaxis] 369 | 370 | 371 | class SaltToTensor: 372 | def __init__(self, coord_channels=False): 373 | self.img2tensor = ImageToTensor(coord_channels) 374 | self.msk2tensor = MaskToTensor() 375 | 376 | def __call__(self, image, mask=None): 377 | if mask is None: 378 | return self.img2tensor(image) 379 | return self.img2tensor(image), self.msk2tensor(mask) 380 | 381 | 382 | class Dummy: 383 | def __call__(self, image, mask=None): 384 | return image, mask 385 | 386 | 387 | class SaltTransform: 388 | def __init__(self, size, train, resize_mode='pad', coord_channels=False): 389 | assert resize_mode in ['scale', 'pad', 'crop'] 390 | self.train = train 391 | 392 | if train: 393 | if resize_mode == 'scale': 394 | resize = Scale(size) 395 | elif resize_mode == 'pad': 396 | resize = OneOf([ 397 | PadToSize(size, random_center=True), 398 | PadToSize(size, random_center=False) 399 | ], prob=1) 400 | elif resize_mode == 'crop': 401 | resize = RandomCrop(size) 402 | else: 403 | raise Exception 404 | self.transform = Compose([ 405 | resize, 406 | UseWithProb(HorizontalFlip(), 0.5), 407 | SaltToTensor(coord_channels) 408 | ]) 409 | else: 410 | if resize_mode == 'scale': 411 | resize = Scale(size) 412 | elif resize_mode == 'pad': 413 | resize = PadToSize(size, random_center=False) 414 | elif resize_mode == 'crop': 415 | resize = CenterCrop(size) 416 | else: 417 | raise Exception 418 | self.transform = Compose([ 419 | resize, 420 | SaltToTensor(coord_channels) 421 | ]) 422 | 423 | def __call__(self, image, mask=None): 424 | if mask is None: 425 | image = self.transform(image) 426 | return image 427 | else: 428 | image, mask = self.transform(image, mask) 429 | return image, mask 430 | --------------------------------------------------------------------------------