spectrumTools.py

"""A set of functions to help load and work with spectral emission, reflectance, and sensitiviy curves"""
import csv
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import savgol_filter

targetWavelengthRange = [420, 650]

#Curve Object
# - startWavelength, endWavelength
# - startYAxisPixel, endYAxisPixel
# - startYAxisRatio, endYAxisRatio
# - Curve Points

def makeCurveObject(curve, wavelengthRange, yAxisPixelRange, yAxisRatioRange):
    """Takes a set of raw curve points and the parameters bounding the curve and returns a curve object"""
    curveObject = {}
    curveObject['curve'] = curve
    curveObject['wavelengthRange'] = wavelengthRange
    curveObject['yAxisPixelRange'] = yAxisPixelRange
    curveObject['yAxisRatioRange'] = yAxisRatioRange

    return curveObject

def copyCurveObject(curve, basisCurveObject):
    """Takes in a set of curve points and copys the parameters from a curve object"""

    curveObject = {}
    curveObject['curve'] = curve
    curveObject['wavelengthRange'] = basisCurveObject['wavelengthRange']
    curveObject['yAxisPixelRange'] = basisCurveObject['yAxisPixelRange']
    curveObject['yAxisRatioRange'] = basisCurveObject['yAxisRatioRange']

    return curveObject

def readCurve(path):
    """
    Read in from a CSV file
        Row 0: [start wavelength, end wavelength]
        Row 1: [y pixel start, y pixel value end] - used to define the y axis range of the plot being read in (i.e. y range in 0 to 100)
        Row 2: [y axis ratio start, y axis ratio end] - used to define the ratio of the y axis being read in (i.e. 0 to 0.55)
        Row 3+: [wavelength, y value]
    """
    with open(path, newline='') as csvFile:
        reader = csv.reader(csvFile, delimiter=',')
        rawFilePairs = [pair for pair in reader]

        wavelengthRange = [int(num) for num in rawFilePairs[0]]
        yAxisPixelRange = [int(num) for num in rawFilePairs[1]]
        yAxisRatioRange = [float(num) for num in rawFilePairs[2]]
        curve = np.asarray([[float(pair[0]), float(pair[1])] for pair in rawFilePairs[3:]])

        return makeCurveObject(curve, wavelengthRange, yAxisPixelRange, yAxisRatioRange)

def readTracedCurves(fileName):
    """Helper function to read in a traced curve - i.e. a curve that was generated by tracing a raster image"""
    return readCurve('curves/tracedCurves/{}.csv'.format(fileName))

def readMeasuredCurves(fileName):
    """Helper function to read in a measured curve - i.e. a curve calculated from an image + spectroscope"""
    return readCurve('curves/measuredCurves/{}.csv'.format(fileName))

def writeCurve(path, curveObject):
    """
    Write to a CSV file
        Row 0: [start wavelength, end wavelength]
        Row 1: [y pixel start, y pixel value end] - used to define the y axis range of the plot being read in (i.e. y range in 0 to 100)
        Row 2: [y axis ratio start, y axis ratio end] - used to define the ratio of the y axis being read in (i.e. 0 to 0.55)
        Row 3+: [wavelength, y value]
    """
    with open(path, 'w', newline='') as csvFile:
        writer = csv.writer(csvFile, delimiter=',', quotechar='|', quoting=csv.QUOTE_MINIMAL)
        writer.writerow(curveObject['wavelengthRange'])
        writer.writerow(curveObject['yAxisPixelRange'])
        writer.writerow(curveObject['yAxisRatioRange'])
        writer.writerows(curveObject['curve'])

def writeMeasuredCurve(fileName, curveObject):
    """Helper function to write a measured curve - i.e. a curve calculated from an image + spectroscope"""
    writeCurve('curves/measuredCurves/{}.csv'.format(fileName), curveObject)

def writeComputedCurve(fileName, curveObject):
    """Helper function to write a computed curve - i.e. a curve calculated individual R, G, B spectra"""
    writeCurve('curves/computedCurves/{}.csv'.format(fileName), curveObject)

def quantizeCurve(curveObject):
    """Interpolate the curve point for each wavelength by simply finding the line between each set of neighboring points"""
    curve = curveObject['curve']

    pointDiffs = curve[1:] - curve[:-1]

    slopes = pointDiffs[:, 1] / pointDiffs[:, 0]
    intercepts = curve[:-1, 1] - (slopes * curve[:-1, 0])
    slopeRanges = np.stack([curve[:-1, 0], curve[1:, 0]], axis=1)

    points = []
    for m, b, lineRange in zip(slopes, intercepts, slopeRanges):
        if (lineRange[0] == lineRange[1]) or (lineRange[0] == (lineRange[1] + 1)):
            xValues = lineRange[0]
        else:
            xValues = np.arange(np.floor(lineRange[0]), np.floor(lineRange[1]))
        yValues = xValues * m + b
        points.append(np.stack([xValues, yValues], axis=1))

    curve = np.concatenate(points, axis=0)
    curveObject['curve'] = curve

    return curveObject

#Smooth after scaling!
def smoothCurve(curveObject):
    """Smooths the curve"""
    curve = curveObject['curve']
    y = savgol_filter(curve[:, 1], 15, 3)
    curveObject['curve'] = np.stack([curve[:, 0], y], axis=1)
    return curveObject

def scaleCurve(curveObject):
    """Scales the x values in the curve to the wavelength range. Scales the y values in the curve to a range of 0-1"""
    curve = curveObject['curve']
    startWavelength, endWavelength = curveObject['wavelengthRange']
    wavelengthRange = endWavelength - startWavelength

    startYAxisPixel, endYAxisPixel = curveObject['yAxisPixelRange']
    yAxisPixelRange = endYAxisPixel - startYAxisPixel

    startYAxisRatio, endYAxisRatio = curveObject['yAxisRatioRange']
    yAxisRatioRange = endYAxisRatio - startYAxisRatio

    maxX, _ = np.max(curve, axis=0)
    minX, _ = np.min(curve, axis=0)

    x = (((curve[:, 0] - minX) / (maxX - minX)) * wavelengthRange) + startWavelength
    y = ((curve[:, 1] / yAxisPixelRange) * yAxisRatioRange) + startYAxisRatio

    curveObject['curve'] = np.stack([x, y], axis=1)
    return curveObject

def cropCurve(curveObject):
    """Crop the curve to the target wavelength range"""
    curve = curveObject['curve']
    start, end = targetWavelengthRange
    startIndex = np.argmax(curve[:, 0] == start)
    endIndex = np.argmax(curve[:, 0] == end)
    curve = curve[startIndex:endIndex, :]
    curveObject['curve'] = curve
    curveObject['wavelengthRange'] = targetWavelengthRange
    return curveObject

def plotCurve(curveObject, marker, show=True):
    """Plots curve"""
    curve = curveObject['curve']
    plt.plot(curve[:, 0], curve[:, 1], marker)

    if show:
        plt.show()

def plotRGBCurves(curveObjects, show=True):
    """Plots RGB curves"""
    redCurve, greenCurve, blueCurve = curveObjects

    plotCurve(redCurve, 'r-', False)
    plotCurve(greenCurve, 'g-', False)
    plotCurve(blueCurve, 'b-', show)

def invertCurve(curveObject):
    """Inverts/flips the curve"""
    curve = curveObject['curve']
    _, endYAxisPixel = curveObject['yAxisPixelRange']
    curve[:, 1] = endYAxisPixel - curve[:, 1]
    curveObject['curve'] = curve
    return curveObject

def getRegionCurveObject(name):
    """Helper function to load and process the region curves"""
    regionObject = readTracedCurves(name)
    correctedOrientation = invertCurve(regionObject)
    scaled = scaleCurve(correctedOrientation)
    quantized = quantizeCurve(scaled)
    smoothed = smoothCurve(quantized)
    cropped = cropCurve(smoothed)
    return cropped

def getEyeCurveObjects():
    """Helper function to load and process the human eye RGB sensitivity curves"""
    redEyeCurve = readTracedCurves('eyeSensitivityRed')
    greenEyeCurve = readTracedCurves('eyeSensitivityGreen')
    blueEyeCurve = readTracedCurves('eyeSensitivityBlue')

    scaledRedEyeCurve = scaleCurve(redEyeCurve)
    scaledGreenEyeCurve = scaleCurve(greenEyeCurve)
    scaledBlueEyeCurve = scaleCurve(blueEyeCurve)

    quantizedRedEyeCurve = quantizeCurve(scaledRedEyeCurve)
    quantizedGreenEyeCurve = quantizeCurve(scaledGreenEyeCurve)
    quantizedBlueEyeCurve = quantizeCurve(scaledBlueEyeCurve)

    smoothedRedEyeCurve = smoothCurve(quantizedRedEyeCurve)
    smoothedGreenEyeCurve = smoothCurve(quantizedGreenEyeCurve)
    smoothedBlueEyeCurve = smoothCurve(quantizedBlueEyeCurve)

    croppedRedEyeCurve = cropCurve(smoothedRedEyeCurve)
    croppedGreenEyeCurve = cropCurve(smoothedGreenEyeCurve)
    croppedBlueEyeCurve = cropCurve(smoothedBlueEyeCurve)

    return [croppedRedEyeCurve, croppedGreenEyeCurve, croppedBlueEyeCurve]

def getMeasuredCurveObjects(name):
    """Helper function to load and process the measured RGB curve values"""
    redObject = readMeasuredCurves('{}_red'.format(name))
    greenObject = readMeasuredCurves('{}_green'.format(name))
    blueObject = readMeasuredCurves('{}_blue'.format(name))

    scaledRed = scaleCurve(redObject)
    scaledGreen = scaleCurve(greenObject)
    scaledBlue = scaleCurve(blueObject)

    quantizedRed = quantizeCurve(scaledRed)
    quantizedGreen = quantizeCurve(scaledGreen)
    quantizedBlue = quantizeCurve(scaledBlue)

    smoothedRed = smoothCurve(quantizedRed)
    smoothedGreen = smoothCurve(quantizedGreen)
    smoothedBlue = smoothCurve(quantizedBlue)

    croppedRed = cropCurve(smoothedRed)
    croppedGreen = cropCurve(smoothedGreen)
    croppedBlue = cropCurve(smoothedBlue)

    return [croppedRed, croppedGreen, croppedBlue]

def getGroundtruthSunlightCurveObject():
    """Helper function to load and process the ground truth sunlight curve """
    sunlightObject = readTracedCurves('sunlight')
    scaled = scaleCurve(sunlightObject)
    quantized = quantizeCurve(scaled)
    smoothed = smoothCurve(quantized)
    cropped = cropCurve(smoothed)
    return cropped

def generateCalibrationCurve(groundTruthSunlightCurveObject, measuredSunlightCurveObject):
    """
    Returns a calibration curve for a digital sensor
        Takes the ground truth sunlight specular emission curve (D65)
        Takes a combined RGB specular sensitivity curve
          Captured through the device to be calibrated of sunlight on a white reference card in early afternoon (Roughly approximate D65)
    """
    groundTruthSunlightCurve = groundTruthSunlightCurveObject['curve']
    measuredSunlightCurve = measuredSunlightCurveObject['curve']

    calibrationCurveX = measuredSunlightCurve[:, 0]
    calibrationCurveY = groundTruthSunlightCurve[:, 1] / measuredSunlightCurve[:, 1]

    calibrationCurve = np.stack([calibrationCurveX, calibrationCurveY], axis=1)

    return copyCurveObject(calibrationCurve, measuredSunlightCurveObject)

def calibrateCurve(curveObject, calibrationCurveObject):
    """Applys a calibration curve to the curve object, returning a new calibrated curve object"""
    calibrationCurve = calibrationCurveObject['curve']
    curve = curveObject['curve']

    calibratedCurve = np.stack([curve[:, 0], (curve[:, 1] * calibrationCurve[:, 1])], axis=1)
    calibratedCurve[:, 1] /= np.max(calibratedCurve[:, 1])

    return copyCurveObject(calibratedCurve, curveObject)

def combineRGBCurves(rgbCurveObjects):
    """Combines individual RGB spectral sensitivite curves into a single combined spectral sensitivity curve"""
    redCurveObject, greenCurveObject, blueCurveObject = rgbCurveObjects

    redCurve = redCurveObject['curve']
    greenCurve = greenCurveObject['curve']
    blueCurve = blueCurveObject['curve']

    combinedCurve = np.copy(redCurve)
    combinedCurve[:, 1] = np.sum([redCurve[:, 1], greenCurve[:, 1], blueCurve[:, 1]], axis=0)

    maxValue = max(combinedCurve[:, 1])

    combinedCurve[:, 1] /= maxValue

    return copyCurveObject(combinedCurve, redCurveObject)

def getLightSourceCurve(name):
    """Returns the curve for a specific light source"""
    rgbCurves = getMeasuredCurveObjects(name)
    combinedCurve = combineRGBCurves(rgbCurves)
    calibratedCurve = calibrateCurve(combinedCurve, d65CalibrationCurve)
    return calibratedCurve

groundTruthSunlight = getGroundtruthSunlightCurveObject()
rgbSunCurves = getMeasuredCurveObjects('Sun')
measuredSunCurve = combineRGBCurves(rgbSunCurves)
d65CalibrationCurve = generateCalibrationCurve(groundTruthSunlight, measuredSunCurve)

#plotCurve(europe1, 'r-', False)
#plotCurve(europe2, 'b-', False)
#plotCurve(europe3, 'g-', False)
#plotCurve(southAsia1, 'r-', False)
#plotCurve(southAsia2, 'b-', False)
#plotCurve(southAsia3, 'g-', False)
#plotCurve(eastAsia1, 'r-', False)
#plotCurve(eastAsia2, 'b-', False)
#plotCurve(eastAsia3, 'g-', False)
#plotCurve(africa1, 'r-', False)
#plotCurve(africa2, 'b-', False)
#plotCurve(africa3, 'g-', False)

#plotCurve(groundTruthSunlight, 'y-', False)
#plotCurve(ledCurve, 'b-')
#plotCurve(skyCurve, 'b--', False)
#plotCurve(incACurve, 'r-', False)
#plotCurve(incBCurve, 'g-', False)
#plotCurve(ipadCurve, 'k-')

#plotRGBCurves(rgbSunCurves, False)
#plotCurve(checkCurve, 'r*', False)
#plotCurve(groundTruthSunlight, 'y-', False)
#plotCurve(combinedSunCurve, 'k-')

#plotRGBCurves(rgbBenQCurves)
#plotRGBCurves(rgbSkyCurves)
#plotRGBCurves(rgbSunCurves)