'''Python example that shows you how to load a segmentation mask from the yeast segmentation
 web tool and draw outlines over a input image'''

import os
import numpy as np
from PIL import Image
import skimage.morphology
import skimage.exposure
import matplotlib.pyplot as plt

# Path of the original raw image
image_name = r"C:\Users\example\Desktop\original_image.png"

# Path of the processed segmentation file
segmentation_name = r"C:\Users\example\Desktop\segmentation.png"

# Load the original raw image and rescale its intensity values
image = np.array(Image.open(image_name))
image = skimage.exposure.rescale_intensity(image.astype(np.float32), out_range=(0, 1))
image = np.round(image * 255).astype(np.uint8)

# Convert the image to an RGB image, if necessary
if len(image.shape) == 3 and image.shape[2] == 3:
    pass
else:
    image = np.expand_dims(image, axis=-1)
    image = np.tile(image, 3)

# Open the segmentation file
seg = np.array(Image.open(segmentation_name))

# Create a raw file to store the outlines
outlines = np.zeros(seg.shape)

# Iterate over each integer in the segmentation and save the outline of each cell onto the outline file
for i in range(1, int(np.max(seg) + 1)):
    tmp = np.zeros(seg.shape)
    tmp[np.where(seg == i)] = 1
    tmp = tmp - skimage.morphology.binary_erosion(tmp)
    outlines += tmp

# Overlay the outlines on the original image in blue
image_outlined = image.copy()
image_outlined[outlines > 0] = (0, 255, 0)

# Display the outline file
plt.imshow(image_outlined)
plt.show()