#
# Color spaces and color matching lab
# James Tompkin
# CSCI 1290
# Brown University
#

import numpy as np
from skimage import color
from PIL import Image
import matplotlib.pyplot as plt

####################################################################################
# Functions
####################################################################################
#
# Use matplotlib to show a 3D scatter plot with subsampling of the input points
# https://matplotlib.org/3.2.1/gallery/mplot3d/scatter3d.html
#
# Input:
# - img: Input image whose intensity values (or 'positions') are the coordinates we wish to plot in a color space
# - colors: The colors associated with those intensity values for visual output to the screen.
# - n: Random number of pixels to plot
# - label_x: Labels for axis.
# - lim_x: Limits for the axis of the color space
def plotColor( positions, colors, n, label_x, label_y, label_z, lim_x, lim_y, lim_z ):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    h,w,c = img.shape

    for i in range(0,n):
        y, x = np.random.randint(h), np.random.randint(w)
        ax.scatter( positions[y,x,0], positions[y,x,1], positions[y,x,2], marker='o', facecolors=colors[y,x,:]/255.0 )

    # e.g., for RGB, this would be 'R', 'G', 'B'
    ax.set_xlabel( label_x )
    ax.set_ylabel( label_y )
    ax.set_zlabel( label_z )

    # e.g., for uint8 RGB, this would be (0,255)
    ax.set_xlim( lim_x )
    ax.set_ylim( lim_y )
    ax.set_zlim( lim_z )

    plt.show()


# Show three images in a column
def showColorTransfer( input, colortarget, transfer ):
    fig = plt.figure()
    ax = fig.add_subplot(311)
    ax.set_axis_off()
    plt.imshow( input.astype(np.uint8) )
    ax = fig.add_subplot(312)
    ax.set_axis_off()
    plt.imshow( colortarget.astype(np.uint8) )
    ax = fig.add_subplot(313)
    ax.set_axis_off()
    plt.imshow( transfer.astype(np.uint8) )
    plt.show()


# You might want to write a convenience function to normalize an image.
# def normalizeI( .... ):

#     return ....


# You might want to write a convenience function to transfer color between two images.
# def colorTransfer( ... ):

#     return ....



####################################################################################
# Script
####################################################################################
#
# Today, we're going to think about three different color spaces, 
# and perform linear transforms on those color spaces to perform color transfer.
# 
# This method of color transfer was designed by Erik Reinhard, Michael Ashikhmin, Bruce Gooch, and Peter Shirley, in 2001.
# It is described in the paper: 'Color Transfer Between Images' 
# https://www.cs.tau.ac.il/~turkel/imagepapers/ColorTransfer.pdf
#
#
# Note: Yes, it's the same person (Erik Reinhard) who designed the global tone mapping operator we use in the HDR project!


###############################
# Part 0: Look at the pictures!
###############################
# Download the paper and have a look at the figures to get a sense of what it is we wish to accomplish.


# Next, download the three images we'll be using today. These are from Figure 1 in the paper.
# - Input image: https://cs.brown.edu/courses/cs129/labs/lab_colormatching/input.png
# - Color target image: https://cs.brown.edu/courses/cs129/labs/lab_colormatching/colortarget.png
# - Reinhard output image: https://cs.brown.edu/courses/cs129/labs/lab_colormatching/reinhard-output.png


#####################
# Code:
# Fix random seed:
np.random.seed(1290)

# Load images
img = np.array(Image.open('input.png'), dtype=np.float32)
ct = np.array(Image.open('colortarget.png'), dtype=np.float32)
# Intended output
img_ct_reinhard = np.array(Image.open('reinhard-output.png'), dtype=np.float32)


###############################################
# Part 1: Visualizing the input RGB color space
###############################################
# Let's look at what data points we're dealing with.
#
# Visualize the RGB points of the input image within an RGB color cube as a scatter plot.
# Then, compare this to the scatter plot of the RGB points of the color target image.
# We've provided a function `plotColor` to draw a scatter plot.
# 
#
# Notes: 
# - Label your axes! 'R', 'G', 'B'
# - Set the correct limits for your axes! For uint8 RGB, this would be (0,255)
#
# - Each pixel is a point - that's a lot of points to draw on a scatter plot!
# - We will subsample the points for viewing using the 'n' parameter, which will randomly select 'n' points to show.
# - Plot with 1000 points first to get a sense of the shape.
# - You will need to lower the number of points for interactive control - James' laptop can do interactive rendering (usable pan-zoom) with 100 points.
#

# Things to note:
# - Within the RGB color cube, the intensity is mixed with the color.
# - The two images have different characteristics within the space - slopes, scales of the 'surfaces'.



#########################################
# Part 2: Color matching in the RGB space
#########################################
# What if we could transfer the color characteristics between two images by mapping the statistics of their pixels within the RGB space?
# Let's translate and scale one set of pixels to have the same statistics as the other set, such that their overall color characteristics match better.
# 
# Steps:
# - First, compute the mean and standard deviation of the input image pixels.
# - Second, compute the mean and standard deviation of the color target image pixels.
# - Third, 'normalizing' or 'standardizing' the input image pixel color distribution by subtracting its mean and dividing by its standard deviation.
# - Fourth, map the input colors into the color target by multiplying by their standard deviation and adding back the mean of the color target pixels.
# 


# Notes:
#   - Do I want a mean/stddev per color channel, or across the whole color space?
#
#   - We've provided a function `showColorTransfer()` to help you see your result.
#   - Make sure to clip your output image, as values > 255 or < 0 will be highlighted by imshow(). RGB values range [0,255]
#
#   - We could scatter plot the colors again to show their match (perhaps with an overlay), but take care:
#   - The 'locations' of the points have now changed, and are in a different limit range.
#   - The colors we want to display on the screen have not - this is what the second argument in `plotColor' is for.



# Questions:
# - What happened to the output?
# - Was it successful?



#########################################
# Part 3: Color matching in the luminance/chrominance space
#########################################
# Let's transform the input image into a color space that is more useful.
# Let's separate out the intensity from the color --- the luminance from the chrominance --- and try it again.

# Step 1: 
# Compute the luminance by applying approximate spectral output display weights to each channel to compute a grayscale luminance image.
# Take the sum of the color channels weighted by the perceptual RGB to gray weights (below).
# Then, visualize the luminance! In plt.imshow(), you will need to add "cmap='gray'" as the image is single channel.

# RGB to Gray weights - approximate 
# https://scikit-image.org/docs/dev/api/skimage.color.html#skimage.color.rgb2gray
rgb2gray_weightr = 0.2125 
rgb2gray_weightg = 0.7154
rgb2gray_weightb = 0.0721


# Step 2: Divide through each color channel by the grayscale to factor out the intensity and leave us with just the color or 'chrominance'.
# The result of this will be three chrominance channels (for red, green, and blue).
# Visualize the chrominance!

# Step 3: Find the mean and std for each image part individually - luminance and chrominance.

# Step 4: Transfer the color by mapping the statistics as before.

# Step 5: Reconstruct the image by multiplying the chroma through by the luminance image again.

# Remember to clip your output as needed.
# Luminance ranges [0,255]


# Step 6: Visualize your output.


# Questions:
# - What happened to the output?
# - Was it successful?
#
# - Note: you might see 'blocky artifacts' in the output due to compression, and your output will not look as clean as `reinhard-output.png`
#   This is expected; the image lost some quality when it was published in the PDF.


################################################
# Part 4: Color transfer in perceptual Lab space
################################################
# Now convert into the perceptual Lab space and try again
# We can use a prebuilt function for this: skimage.color.rgb2lab()
# 

# Step 1: Convert images into Lab space
# Step 2: Apply color transfer technique by aligning statistics of pixels in color space
# Step 3: Convert back to RGB space

# Notes:
#  - skimage.color.rgb2lab() expects inputs in the range [0,1], not [0,255]
#  - Make sure to clamp your output again - L channel ranges from [0,100], a,b channels range[-100,100]


# Questions:
# - What happened to the output?
# - Was it successful?


#####################################################################
# Part 5: Across the three techniques, compare the outputs visually!
#####################################################################
# We've used three color spaces for color transfer: RGB, luminance/chrominance (similar to YUV), and Lab.
# Now, let's compare their outputs.
# You can reuse `showColorTranfer()` to show all three.


# Questions:
#  - Which do you prefer?
#  - What do you notice about the different color transfer techniques?
#  - What are the causes of the differences?


#########################################
# Part 6: Capture your own images!
#########################################
# Capture your own images and test them with the provided color target image.
# Or, capture new color target images and test them, too.
#
# For complex images, Reinhard and colleagues had to adjust different image areas and apply different weights.
# For instance, masking out the sky in one scene from the landscape, and applying two separate color transfers.
# 
# What can you create?


#########################################
# Part 7: Upload your image to Gradescope
#########################################
# Put your image into a PDF and upload it to Gradescope.
# Lab = done.

# Operating in the right color space can have a huge difference on the result - separating intensity from color is critical.