{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"


Stain Normalization 👀: StainGAN, StainNet 🔬

\n \n
\n\n
CREATED BY: NGHI HUYNH
","metadata":{}},{"cell_type":"markdown","source":"

\n


CONTENTS

\n\n\n

0    IMPORTS

\n\n---\n\n

1    STAIN NORMALIZATION TECHNIQUES

\n\n---\n\n

2    EVALUATION

\n\n---\n\n

3    DATASET CREATION

\n\n---","metadata":{}},{"cell_type":"markdown","source":"
\n


If you find this notebook useful, do give me an upvote, it motivates me a lot.

This notebook is still a work in progress. Keep checking for further developments!😊

","metadata":{}},{"cell_type":"markdown","source":"\n\n


 0. IMPORTS ⤒

","metadata":{}},{"cell_type":"code","source":"!pip install staintools\n!pip install spams","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:11:31.746760Z","iopub.execute_input":"2022-08-13T19:11:31.747552Z","iopub.status.idle":"2022-08-13T19:13:31.904336Z","shell.execute_reply.started":"2022-08-13T19:11:31.747513Z","shell.execute_reply":"2022-08-13T19:13:31.903039Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nimport re\nimport gc\nimport time\nimport os\nimport zipfile\nimport cv2\nimport staintools\nimport pandas as pd\n\nfrom glob import glob\nfrom tqdm.notebook import tqdm\nimport torch\nimport numpy as np\nimport matplotlib.pyplot as plt \nfrom skimage import measure\nfrom PIL import Image \nimport matplotlib.pyplot as plt\nimport tifffile as tiff\n\nfrom skimage.metrics import structural_similarity as ssim, peak_signal_noise_ratio as psnr\n%matplotlib inline","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:15:23.948921Z","iopub.execute_input":"2022-08-13T19:15:23.949558Z","iopub.status.idle":"2022-08-13T19:15:26.812530Z","shell.execute_reply.started":"2022-08-13T19:15:23.949519Z","shell.execute_reply":"2022-08-13T19:15:26.811559Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sys.path.append('../input/StainNet/')","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:15:26.814324Z","iopub.execute_input":"2022-08-13T19:15:26.814973Z","iopub.status.idle":"2022-08-13T19:15:26.820187Z","shell.execute_reply.started":"2022-08-13T19:15:26.814936Z","shell.execute_reply":"2022-08-13T19:15:26.819199Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from models import StainNet, ResnetGenerator","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:15:27.562819Z","iopub.execute_input":"2022-08-13T19:15:27.563391Z","iopub.status.idle":"2022-08-13T19:15:27.809821Z","shell.execute_reply.started":"2022-08-13T19:15:27.563356Z","shell.execute_reply":"2022-08-13T19:15:27.808460Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n\n


 1. STAIN NORMALIZATION TECHNIQUES ⤒

","metadata":{}},{"cell_type":"markdown","source":"Following my discussion thread on [Stain Normalization Techniques](https://www.kaggle.com/competitions/hubmap-organ-segmentation/discussion/338525), I created this notebook to explore these techniques in more detail.\n\nHere are a few **conventional** stain normalization techniques developed by **Reinhard**, **Macenko**, and **Vahadane**. [[*reference: A review of stain removal*](https://iopscience.iop.org/article/10.1088/1742-6596/1362/1/012108/pdf)]\n\n* **Reinhard**: based on color transfer between a standard image and color varied image using mean and variance of both the images. Then. alter the source image to the target image\n\n* **Macenko**: find particular stain vectors for each image based on the colors that are present in the image\n\n* **Vahadane**: decompose the image into stain density map that are sparse and non-negative. Then, the stain density maps are combined on the basis of stain color of a pathologist preferred target image. Thus, altering only its color and preserving the structure.\n\nHere are some previous works based on **deep-learning**-based methods primarily applying generative adversarial networks (GANs) as well.\n\n* **StainGAN**: based on CycleGAN to transfer the stain style. Limitation: complex and might have a risk of introducing some artifacts [[*reference: Cycle-consistent adversial networks*](https://arxiv.org/abs/1703.10593)]\n\n* **StainNet**: uses StainGAN as the teacher network, to learn the color mapping by distillation learning. [[*reference: StainNet-a fast and robust network*](https://www.frontiersin.org/articles/10.3389/fmed.2021.746307/full)] \n","metadata":{}},{"cell_type":"markdown","source":"
\n


Even though stain normalization has some proven benefits for training our models, experts in this field recommend that it's not always useful. Since we want to build a model being robust to color, it's better to do color augmentation instead of normalizing them. However, you can still use it as long as it's reasonable.😊

","metadata":{}},{"cell_type":"markdown","source":"## *Helper Functions*","metadata":{}},{"cell_type":"code","source":"# https://github.com/khtao/StainNet/blob/master/demo.ipynb\ndef norm(image):\n image = np.array(image).astype(np.float32)\n image = image.transpose((2, 0, 1))\n image = ((image / 255) - 0.5) / 0.5\n image=image[np.newaxis, ...]\n image=torch.from_numpy(image)\n return image\n\ndef un_norm(image):\n image = image.cpu().detach().numpy()[0]\n image = ((image * 0.5 + 0.5) * 255).astype(np.uint8).transpose((1,2,0))\n return image","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:15:32.007222Z","iopub.execute_input":"2022-08-13T19:15:32.007595Z","iopub.status.idle":"2022-08-13T19:15:32.014948Z","shell.execute_reply.started":"2022-08-13T19:15:32.007562Z","shell.execute_reply":"2022-08-13T19:15:32.013647Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def highlight(row):\n df = lambda x: ['background: #CCCCFF' if x.name in row\n else '' for i in x]\n return df","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:15:34.460996Z","iopub.execute_input":"2022-08-13T19:15:34.461425Z","iopub.status.idle":"2022-08-13T19:15:34.467327Z","shell.execute_reply.started":"2022-08-13T19:15:34.461389Z","shell.execute_reply":"2022-08-13T19:15:34.466270Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## *Source & Target images*","metadata":{}},{"cell_type":"code","source":"source = cv2.imread('../input/hubmap-256x256/train/10044_0004.png', cv2.COLOR_BGR2RGB )\ntarget = tiff.imread('../input/hubmap-organ-segmentation/test_images/10078.tiff')","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:15:36.733835Z","iopub.execute_input":"2022-08-13T19:15:36.734401Z","iopub.status.idle":"2022-08-13T19:15:36.965854Z","shell.execute_reply.started":"2022-08-13T19:15:36.734366Z","shell.execute_reply":"2022-08-13T19:15:36.964838Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(16,9))\nplt.subplot(121)\nplt.title(f\"Source\\n{source.shape}\")\nplt.imshow(source)\nplt.subplot(122)\nplt.title(f'Target\\n{target.shape}')\nplt.imshow(target)","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:15:38.891173Z","iopub.execute_input":"2022-08-13T19:15:38.891558Z","iopub.status.idle":"2022-08-13T19:15:39.912421Z","shell.execute_reply.started":"2022-08-13T19:15:38.891525Z","shell.execute_reply":"2022-08-13T19:15:39.910766Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## *Conventional Techniques*","metadata":{}},{"cell_type":"markdown","source":"### Reinhard","metadata":{}},{"cell_type":"code","source":"# run reinhard normalization\nnormalizer = staintools.ReinhardColorNormalizer()\nnormalizer.fit(np.array(target))\nreinhard_normalized = normalizer.transform(np.array(source))\nplt.imshow(reinhard_normalized)\nplt.title('Reinhard')","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:15:42.945536Z","iopub.execute_input":"2022-08-13T19:15:42.946160Z","iopub.status.idle":"2022-08-13T19:15:43.434521Z","shell.execute_reply.started":"2022-08-13T19:15:42.946122Z","shell.execute_reply":"2022-08-13T19:15:43.433294Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Vahadane","metadata":{}},{"cell_type":"code","source":"# run vahadane normalization\nnormalizer = staintools.StainNormalizer(method=\"vahadane\")\nnormalizer.fit(np.array(target))\nvahadane_normalized = normalizer.transform(np.array(source))\nplt.imshow(vahadane_normalized)\nplt.title('Vahadane')","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:15:45.597212Z","iopub.execute_input":"2022-08-13T19:15:45.597814Z","iopub.status.idle":"2022-08-13T19:15:53.614428Z","shell.execute_reply.started":"2022-08-13T19:15:45.597778Z","shell.execute_reply":"2022-08-13T19:15:53.613367Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## *Deep Learning Techniques*","metadata":{}},{"cell_type":"markdown","source":"### StainGAN","metadata":{}},{"cell_type":"code","source":"# load pretrained StainGAN\nmodel_GAN = ResnetGenerator(3, 3, ngf=64, norm_layer=torch.nn.InstanceNorm2d, n_blocks=9).cuda()\nmodel_GAN.load_state_dict(torch.load('../input/StainNet/checkpoints/aligned_histopathology_dataset/latest_net_G_A.pth'))","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:15:53.619245Z","iopub.execute_input":"2022-08-13T19:15:53.621514Z","iopub.status.idle":"2022-08-13T19:15:57.868838Z","shell.execute_reply.started":"2022-08-13T19:15:53.621475Z","shell.execute_reply":"2022-08-13T19:15:57.867666Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_GAN.eval()\nwith torch.no_grad():\n img_gan=model_GAN(norm(source).cuda())\n img_gan=un_norm(img_gan)\n plt.imshow(img_gan)\n plt.title('StainGAN')","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:15:57.870740Z","iopub.execute_input":"2022-08-13T19:15:57.871134Z","iopub.status.idle":"2022-08-13T19:16:03.433391Z","shell.execute_reply.started":"2022-08-13T19:15:57.871093Z","shell.execute_reply":"2022-08-13T19:16:03.432493Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## StainNet","metadata":{}},{"cell_type":"code","source":"# load pretrained StainNet\nmodel_Net = StainNet().cuda()\nmodel_Net.load_state_dict(torch.load(\"../input/StainNet/checkpoints/aligned_histopathology_dataset/StainNet-Public_layer3_ch32.pth\"))","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:16:03.437170Z","iopub.execute_input":"2022-08-13T19:16:03.439912Z","iopub.status.idle":"2022-08-13T19:16:03.459788Z","shell.execute_reply.started":"2022-08-13T19:16:03.439874Z","shell.execute_reply":"2022-08-13T19:16:03.458761Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_Net.eval()\nwith torch.no_grad():\n img_net=model_Net(norm(source).cuda())\n img_net=un_norm(img_net)\n plt.imshow(img_net)\n plt.title('StainNet')","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:16:03.463804Z","iopub.execute_input":"2022-08-13T19:16:03.464216Z","iopub.status.idle":"2022-08-13T19:16:03.704642Z","shell.execute_reply.started":"2022-08-13T19:16:03.464181Z","shell.execute_reply":"2022-08-13T19:16:03.703678Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"type(img_net)","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:16:03.708757Z","iopub.execute_input":"2022-08-13T19:16:03.709662Z","iopub.status.idle":"2022-08-13T19:16:03.717307Z","shell.execute_reply.started":"2022-08-13T19:16:03.709575Z","shell.execute_reply":"2022-08-13T19:16:03.716017Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n\n


 2. EVALUATION ⤒

","metadata":{}},{"cell_type":"markdown","source":"## *Visual Comparison*","metadata":{}},{"cell_type":"code","source":"normalized_imgs = {'Source': source, 'Target': target, \n 'Reinhard': reinhard_normalized, 'Vahadane': vahadane_normalized, \n 'StainGan': img_gan, 'StainNet': img_net}","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:16:03.719276Z","iopub.execute_input":"2022-08-13T19:16:03.719941Z","iopub.status.idle":"2022-08-13T19:16:03.726436Z","shell.execute_reply.started":"2022-08-13T19:16:03.719906Z","shell.execute_reply":"2022-08-13T19:16:03.724920Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(16,16))\nfor i, (method, img) in enumerate(normalized_imgs.items()):\n plt.subplot(1, 6, i+1)\n plt.title(method)\n plt.imshow(img)\nplt.tight_layout()\n ","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:16:03.728202Z","iopub.execute_input":"2022-08-13T19:16:03.729084Z","iopub.status.idle":"2022-08-13T19:16:05.472235Z","shell.execute_reply.started":"2022-08-13T19:16:03.729035Z","shell.execute_reply":"2022-08-13T19:16:05.463762Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## *Evaluation Metrics*\n\nThere are 2 types of similarity metrics to evaluate the performance of different methods:\n\n* **[Structural Similarity Index (SSIM)](https://scikit-image.org/docs/dev/auto_examples/transform/plot_ssim.html)**\n* **[Peak Signal-to-Noise Ratio (PSNR)](https://scikit-image.org/docs/dev/api/skimage.metrics.html?highlight=structural_similarity#skimage.metrics.peak_signal_noise_ratio)**\n\nThe SSIM Source and PSNR Source are used to evaluate the similarity between the normalized image and the source image. The SSIM Source is used to measure the preservation of the source image texture information. ","metadata":{}},{"cell_type":"code","source":"methods = []\nssim_results = []\npsnr_results = []\nfor i, (method, img) in enumerate(normalized_imgs.items()):\n if i > 1:\n methods.append(method)\n ssim_results.append(ssim(source, img, data_range=img.max() - img.min(), channel_axis=2))\n psnr_results.append(psnr(source, img, data_range=img.max() - img.min()))\n","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:16:05.924176Z","iopub.execute_input":"2022-08-13T19:16:05.924755Z","iopub.status.idle":"2022-08-13T19:16:06.044463Z","shell.execute_reply.started":"2022-08-13T19:16:05.924690Z","shell.execute_reply":"2022-08-13T19:16:06.043276Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"evaluation_table = pd.DataFrame(list(zip(methods, ssim_results, psnr_results)), \n columns=['Methods', 'SSIM source', 'PSNR source'])\nevaluation_table.style.hide_index().apply(highlight([1,3]), axis=1)","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:16:08.034126Z","iopub.execute_input":"2022-08-13T19:16:08.034744Z","iopub.status.idle":"2022-08-13T19:16:08.137977Z","shell.execute_reply.started":"2022-08-13T19:16:08.034673Z","shell.execute_reply":"2022-08-13T19:16:08.136911Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"=> **StainNet** outperforms other techniques. However, if we want to use a conventional technique instead of a deep-learning-based technique then **Vahadane** is a good to go one.","metadata":{}},{"cell_type":"markdown","source":"\n\n


 3. DATASET CREATION ⤒

","metadata":{}},{"cell_type":"code","source":"PATH = '../input/hubmap-256x256/train'\nOUT_TRAIN = './normalized.zip'","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:16:11.022627Z","iopub.execute_input":"2022-08-13T19:16:11.023013Z","iopub.status.idle":"2022-08-13T19:16:11.027924Z","shell.execute_reply.started":"2022-08-13T19:16:11.022979Z","shell.execute_reply":"2022-08-13T19:16:11.026836Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with zipfile.ZipFile(OUT_TRAIN, 'w') as img_out:\n for _,_,files in os.walk(PATH):\n for file_name in files:\n with torch.no_grad():\n source = Image.open(PATH+'/'+file_name) \n image_net=model_Net(norm(source).cuda())\n image_net=un_norm(image_net)\n im = cv2.imencode('.png',image_net)[1]\n img_out.writestr(f'{file_name}.png', im) ","metadata":{"execution":{"iopub.status.busy":"2022-08-13T19:16:13.360197Z","iopub.execute_input":"2022-08-13T19:16:13.360573Z","iopub.status.idle":"2022-08-13T19:16:46.818921Z","shell.execute_reply.started":"2022-08-13T19:16:13.360532Z","shell.execute_reply":"2022-08-13T19:16:46.817951Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}