Last active
November 28, 2016 22:30
-
-
Save JosephCatrambone/3bb3a49b454bcd52a68e to your computer and use it in GitHub Desktop.
A simple tool to predict the color of a fuse in a given slot.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from __future__ import division | |
import sys | |
import math | |
from PIL import Image | |
def RGBToHSL(rgb): | |
"""Given an RGB tuple in the range of 0-255, returns HSL from 0-1.""" | |
r = rgb[0]/255.0 | |
g = rgb[1]/255.0 | |
b = rgb[2]/255.0 | |
# Calculate luminance | |
low = min((r,g,b)) | |
high = max((r,g,b)) | |
luminance = float(low+high)/2.0 | |
chroma = float(high-low) | |
# Calculate saturation | |
saturation = 0 | |
if chroma != 0: | |
#saturation = chroma / (1.0 - math.fabs(2.0*luminance - 1)) | |
if luminance > 0.5: | |
saturation = chroma/float(2.0-high-low) | |
elif luminance <= 0.5: | |
saturation = chroma/float(high+low) | |
# Calculate hue | |
hue = 0 | |
if saturation != 0: | |
if high == r: | |
hue = (g-b)/chroma # Normally mod6, but we handle hits in hue below. | |
elif high == g: # Green strongest | |
hue = 2.0 + ((b-r)/chroma) | |
else: #high == b: # Blue strongest | |
hue = 4.0 + ((r-g)/chroma) | |
# Normally we multiply by 60 to normalize hue to 360 degrees. | |
# We will do that, make sure it's in the 0-360 range, then map it back to [0,1] | |
hue *= 60 | |
hue = (hue+360) % 360 # Prevent below-zero and above 360 values. | |
hue /= 360 | |
return hue, saturation, luminance | |
class Rectangle(object): | |
def __init__(self, x, y, width, height): | |
self.x = x | |
self.y = y | |
self.width = width | |
self.height = height | |
def inside(self, x, y): | |
if x < self.x: | |
return False | |
if x > self.x+self.width: | |
return False | |
if y < self.y: | |
return False | |
if y > self.y+self.height: | |
return False | |
return True | |
def get_hue_confidence_in_rect(image, rect, hue_buckets=8): | |
""" | |
Given an image and a rectangle, | |
return an array of numbers (where the array is the length of the number of hues), | |
such that the value is the confidence for the given hue. | |
For example, if we said hue_buckets = 3, we would have only three possible hues (red, green, blue), and we passed in a mostly blue picture, we'd get: | |
[0.0, 0.2, 0.8] -> High blue confidence. | |
Or if we passed in a rainbow picture, we'd get | |
[0.35, 0.33, 0.32] -> No certainty. High confusion. | |
If we said hue_buckets = 5, we'd get arrays of size five, roughly [red, orange, yellow, green, blue]. | |
""" | |
buckets = [0.0]*hue_buckets | |
for y in range(rect.y, rect.y+rect.height): | |
for x in range(rect.x, rect.x+rect.width): | |
pixel = image.getpixel((x, y)) | |
# Convert to HSL. | |
hue, saturation, luminance = RGBToHSL(pixel) | |
# Hue is the pure color | |
# Saturation is 0 at grey and 1 at full color | |
# Luminance/Lightness is 0 at black and 1 at pure white | |
# From this, a saturation of 1 means we're more confident in our color selection bin. | |
# It also means a luminance of 0.5 is the most confident in our color selection. | |
# To explain this a bit better, at a saturation of 0.5 (pure color), we evaluate (1.0 - (2.0*0)) -> 1.0 | |
# At 0 (black, no color) and 1 (white, no color), we get (1.0 - (2.0*0.5)) -> 0.0 | |
confidence = luminance * (1.0-(2.0*math.fabs(0.5-saturation))) | |
# Now we have to map our hue to a bucket. | |
bucket = int(hue*(hue_buckets)) | |
# Now add to our bucket the confidence. | |
buckets[bucket] += confidence | |
return buckets | |
def main(filename): | |
# Load an image | |
img = Image.open(filename) | |
# Hard code the eight rectangles. Assume alignment. | |
rects = list() | |
y1 = 265 | |
y2 = 514 | |
width = 85 | |
height = 235 | |
rects.append(Rectangle(144, y1, width, height)) | |
rects.append(Rectangle(289, y1, width, height)) | |
rects.append(Rectangle(419, y1, width, height)) | |
rects.append(Rectangle(562, y1, width, height)) | |
rects.append(Rectangle(833, y1, width, height)) | |
rects.append(Rectangle(972, y1, width, height)) | |
rects.append(Rectangle(144, y2, width, height)) | |
rects.append(Rectangle(289, y2, width, height)) | |
rects.append(Rectangle(419, y2, width, height)) | |
rects.append(Rectangle(562, y2, width, height)) | |
rects.append(Rectangle(833, y2, width, height)) | |
rects.append(Rectangle(972, y2, width, height)) | |
# Finally, get the colors inside each rectangle. | |
# Remember the color wheel is circular, so we end with and start with red. | |
colors = ['red', 'orange', 'yellow', 'green', 'blue', 'indigo', 'violet', 'red'] | |
for fuse_index, rect in enumerate(rects): | |
dist = get_hue_confidence_in_rect(img, rect, hue_buckets=len(colors)) | |
# We're just selecting argmax for now, but if you calculate the variance of dist, you can also get certainty. | |
max_index = -1 | |
max_value = -1 | |
for index, value in enumerate(dist): | |
if value > max_value: | |
max_index = index | |
max_value = value | |
print("Fuse {} color: {}".format(fuse_index, colors[max_index])) | |
if __name__=="__main__": | |
main(sys.argv[1]) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment