#!/usr/bin/env python
# coding: utf-8

# In[1]:


# Install itk-spcn and itkwidgets, if necessary
import sys
necessary = False
if necessary:
    get_ipython().system('{sys.executable} -m pip install itk-spcn')
    get_ipython().system('{sys.executable} -m pip install itkwidgets')


# In[2]:


# Import needed packages
from urllib.request import urlretrieve
import os

# Import itk, which includes itk-spcn.
import itk
from itkwidgets import view


# In[3]:


# Fetch input images, if we don have them already.
input_image_filename = 'Easy1.png'
input_image_url = 'https://data.kitware.com/api/v1/file/576ad39b8d777f1ecd6702f2/download'
if not os.path.exists(input_image_filename):
    urlretrieve(input_image_url, input_image_filename)

reference_image_filename = 'Hard.png'
reference_image_url = 'https://data.kitware.com/api/v1/file/57718cc48d777f1ecd8a883f/download'
if not os.path.exists(reference_image_filename):
    urlretrieve(reference_image_url, reference_image_filename)
    
output_image_filename = 'HardWithEasy1Colors.png'

# The pixels are RGB triplets of unsigned char.  The images are 2 dimensional.
PixelType = itk.RGBPixel[itk.UC]
ImageType = itk.Image[PixelType, 2]


# In[4]:


# Invoke the functional, eager interface for ITK
input_image = itk.imread(input_image_filename, PixelType)
reference_image = itk.imread(reference_image_filename, PixelType)
eager_normalized_image = itk.structure_preserving_color_normalization_filter(
    input_image,
    reference_image,
    color_index_suppressed_by_hematoxylin=0,
    color_index_suppressed_by_eosin=1)
itk.imwrite(eager_normalized_image, output_image_filename)


# In[5]:


view(input_image)


# In[6]:


view(reference_image)


# In[7]:


view(eager_normalized_image)


# In[8]:


# Alternatively, invoke the ITK pipeline
input_reader = itk.ImageFileReader[ImageType].New(FileName=input_image_filename)
reference_reader = itk.ImageFileReader[ImageType].New(FileName=reference_image_filename)
spcn_filter = itk.StructurePreservingColorNormalizationFilter.New(Input=input_reader.GetOutput())
spcn_filter.SetColorIndexSuppressedByHematoxylin(0)
spcn_filter.SetColorIndexSuppressedByEosin(1)
spcn_filter.SetInput(0, input_reader.GetOutput())
spcn_filter.SetInput(1, reference_reader.GetOutput())
output_writer = itk.ImageFileWriter.New(spcn_filter.GetOutput())
output_writer.SetInput(spcn_filter.GetOutput())
output_writer.SetFileName(output_image_filename)
output_writer.Write()


# In[9]:


view(output_writer.GetInput())


# In[ ]: