dominantcolors.ipynb

{
 "metadata": {
  "name": "Finding Colors in Images"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": "### Let's find the dominant colors in images\n\n![](http://media.charlesleifer.com/blog/photos/thumbnails/akira_650x650.jpg)"
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# we will use the python imaging library to read color info from the image\nfrom PIL import Image",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# let's start by pulling down the image data\nimport urllib2\nfh = urllib2.urlopen('http://media.charlesleifer.com/blog/photos/thumbnails/akira_650x650.jpg')\nimg_data = fh.read()\nfh.close()",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# what does img_data look like?  here are the first 10 bytes -- looks like a header and some null bytes\nprint img_data[:10]",
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": "\ufffd\ufffd\ufffd\ufffd\u0000\u0010JFIF\n"
      }
     ],
     "prompt_number": 3
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# let's load up this image data\nfrom StringIO import StringIO\nimg_buf = StringIO(img_data)\nimg = Image.open(img_buf)",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 4
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# what is img? it should be a jpeg image file\nprint img",
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": "<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=650x334 at 0x9B9CFCC>\n"
      }
     ],
     "prompt_number": 5
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": "## What happened\n\nWe just pulled down some image data over the wire and created an Image object.  This Image object gives us a nice, fast way to extract things like raw pixel data, which we will use to determine dominant colors.\n\nNow that we have the image data, we will resize the image down to 200px on a side -- this makes calculations faster since we have less pixels to count."
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# let's resize the image in C, this will make calculations faster and we won't lose much accuracy\nimg.thumbnail((200, 200))",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 6
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# let's load up some modules which will be useful while we're extracting color info and clustering\nfrom collections import namedtuple\nimport random",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 7
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": "We want a nice way to represent the various color points in the image.  I chose to use ``namedtuple``, which has lower memory overhead than a python class.  The point class will store coordinate data, the number of dimensions (always 3 in this case), and a count associated with the point.\n\nThe clusters will be a collection of points, and have the additional property of a \"center\"."
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# these classes will represent the data we extract -- I use namedtuples as they have lower memory overhead than full classes\nPoint = namedtuple('Point', ('coords', 'n', 'ct'))\nCluster = namedtuple('Cluster', ('points', 'center', 'n'))",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 8
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# let's extract all the color points from the image -- the red/green/blue channels will be treated as points in a 3-dimensional space\ndef get_points(img):\n    points = []\n    w, h = img.size\n    for count, color in img.getcolors(w * h):\n        points.append(Point(color, 3, count))\n    return points",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 9
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "img_points = get_points(img)",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 10
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# when we're clustering we will need a way to find the distance between two points\ndef point_distance(p1, p2):\n    return sum([\n        (p1.coords[i] - p2.coords[i]) ** 2 for i in range(p1.n)\n    ])",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 11
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# we also need a way to calculate the center when given a cluster of points -- this is done\n# by taking the average of the points across all dimensions\ndef calculate_center(points, n):\n    vals = [0.0 for i in range(n)]\n    plen = 0\n    for p in points:\n        plen += p.ct\n        for i in range(n):\n            vals[i] += (p.coords[i] * p.ct)\n    return Point([(v / plen) for v in vals], n, 1)",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 12
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": "### What have we done so far\n\nSo far we have extracted the points from the image and created a few helper functions for things like calculating the center of a cluster of points and calculating the distance between two points.\n\n### The next step -- running the algorithm\n\nThe code can be found in the next cell, the algorithm is ``k-means``\n\nOur goal is to find where the points tend to form \u201cclumps\u201d. Since we want to group the numbers into k clusters, we\u2019ll pick k points randomly from the data to use as the initial \u201cclusters\u201d.\n\nWe\u2019ll iterate over every point in the data and calculate its distance to each of the k clusters. Find the nearest cluster and associate that point with the cluster. When you\u2019ve iterated over all the points they should all be assigned to one of the clusters. Now, for each cluster recalculate its center by averaging the distances of all the associated points and start over.\n\nWhen the centers stop moving very much we can stop looping.  To find the dominant colors, simply take the centers of the clusters!"
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# finally, here is our algorithm -- 'kmeans'\ndef kmeans(points, k, min_diff):\n    clusters = [Cluster([p], p, p.n) for p in random.sample(points, k)]\n \n    while 1:\n        plists = [[] for i in range(k)]\n \n        for p in points:\n            smallest_distance = float('Inf')\n            for i in range(k):\n                distance = point_distance(p, clusters[i].center)\n                if distance < smallest_distance:\n                    smallest_distance = distance\n                    idx = i\n            plists[idx].append(p)\n \n        diff = 0\n        for i in range(k):\n            old = clusters[i]\n            center = calculate_center(plists[i], old.n)\n            new = Cluster(plists[i], center, old.n)\n            clusters[i] = new\n            diff = max(diff, point_distance(old.center, new.center))\n \n        if diff < min_diff:\n            break\n \n    return clusters",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 13
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "print 'Calculating clusters -- this may take a few seconds'\nclusters = kmeans(img_points, 3, 1) # run k-means on the color points, calculating 3 clusters (3 dominant colors), and stopping when our clusters move < 1 unit\nrgbs = [map(int, c.center.coords) for c in clusters]\nprint 'Done'",
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": "Calculating clusters -- this may take a few seconds\nDone"
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": "\n"
      }
     ],
     "prompt_number": 14
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "print rgbs",
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": "[[36, 24, 33], [86, 100, 112], [192, 87, 42]]\n"
      }
     ],
     "prompt_number": 15
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# let's create a function to convert RGBs into hex color code\nrtoh = lambda rgb: '#%s' % ''.join(('%02x' % p for p in rgb))",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 16
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "color_codes = map(rtoh, rgbs)\nprint color_codes",
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": "['#241821', '#566470', '#c0572a']\n"
      }
     ],
     "prompt_number": 17
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "# now, let's display those colors using HTML\nfrom IPython.core.display import HTML",
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 18
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "HTML('<div style=\"width: 40px; height: 40px; background-color: %s\">&nbsp;</div>' % color_codes[0])",
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "html": "<div style=\"width: 40px; height: 40px; background-color: #241821\">&nbsp;</div>",
       "output_type": "pyout",
       "prompt_number": 19,
       "text": "<IPython.core.display.HTML at 0x9c580ec>"
      }
     ],
     "prompt_number": 19
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "HTML('<div style=\"width: 40px; height: 40px; background-color: %s\">&nbsp;</div>' % color_codes[1])",
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "html": "<div style=\"width: 40px; height: 40px; background-color: #566470\">&nbsp;</div>",
       "output_type": "pyout",
       "prompt_number": 20,
       "text": "<IPython.core.display.HTML at 0x9c583ac>"
      }
     ],
     "prompt_number": 20
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "HTML('<div style=\"width: 40px; height: 40px; background-color: %s\">&nbsp;</div>' % color_codes[2])",
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "html": "<div style=\"width: 40px; height: 40px; background-color: #c0572a\">&nbsp;</div>",
       "output_type": "pyout",
       "prompt_number": 21,
       "text": "<IPython.core.display.HTML at 0x9c5842c>"
      }
     ],
     "prompt_number": 21
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": "",
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}