{{ message }}

Instantly share code, notes, and snippets.

# JosephCatrambone/fuse_color.py

Last active Nov 28, 2016
A simple tool to predict the color of a fuse in a given slot.
 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])