Median Cut Light Proble Sampling

sampling.py

#!/usr/bin/env python2

import cv2 as cv
import numpy as np
import sys

SIZE = (640, 480)


def m(p, q, f):
  """Compute the pq moment."""

  if p == 0 and q == 0:
    return np.sum(f)

  if p == 1:
    sum = 0
    for index, e in np.ndenumerate(f):
      x, y = index
      sum += x * e
    return sum

  if q == 1:
    sum = 0
    for index, e in np.ndenumerate(f):
      x, y = index
      sum += y * e
    return sum

  return 0


def cut(rgb, gray, maxdepth=7):
  """Runs the median cut algorithm."""

  # Compute the cumulative sums.
  sums = np.cumsum(gray, axis=0, dtype=np.uint)
  sums = np.cumsum(sums, axis=1, dtype=np.uint)

  # Clone the rgb image.
  out = np.copy(rgb)

  def cut_rect(x0, y0, x1, y1, depth=0):
    """Recursively splits a rectangle into areas of equal sum."""

    if depth > maxdepth:
      sliceRGB = rgb[x0:x1, y0:y1]
      sliceGray = gray[x0:x1, y0:y1]

      # Find the centroid.
      m00 = m(0, 0, sliceGray)
      cx = x0 + int(m(1, 0, sliceGray) / m00)
      cy = y0 + int(m(0, 1, sliceGray) / m00)

      # Draw a circle with the extracted diffuse colour.
      diff = np.mean(sliceRGB, axis=(0, 1))
      colour = (int(diff[0]), int(diff[1]), int(diff[2]))

      cv.circle(out, (cy, cx), 4, (0, 0, 0))
      cv.circle(out, (cy, cx), 3, colour, -1)
      return

    total = sums[x1, y1] + sums[x0, y0] - sums[x0, y1] - sums[x1, y0]

    if x1 - x0 > y1 - y0:
      # Cut along y.
      best_x = x0
      best_d = sys.maxint
      for x in range(x0, x1 + 1):
        sum = int(sums[ x, y1] + sums[x0, y0] - sums[x0, y1] - sums[ x, y0])
        if abs(2 * sum - total) < best_d:
          best_d = abs(2 * sum - total)
          best_x = x
      cv.line(out, (y0, best_x), (y1, best_x), (0, 0, 255))
      cut_rect(x0, y0, best_x, y1, depth + 1)
      cut_rect(best_x, y0, x1, y1, depth + 1)
    else:
      # Cut along x.
      best_y = y0
      best_d = sys.maxint
      for y in range(y0, y1 + 1):
        sum = int(sums[x1,  y] + sums[x0, y0] - sums[x0,  y] - sums[x1, y0])
        if abs(2 * sum - total) < best_d:
          best_d = abs(2 * sum - total)
          best_y = y

      cv.line(out, (best_y, x0), (best_y, x1), (0, 0, 255))
      cut_rect(x0, y0, x1, best_y, depth + 1)
      cut_rect(x0, best_y, x1, y1, depth + 1)


  # Start the algorithm, covering the entire image.
  cut_rect(0, 0, gray.shape[0] - 1, gray.shape[1] - 1)
  return out

def main(file):
  """Entry point of the script."""

  image = cv.resize(cv.imread(file, cv.IMREAD_COLOR), SIZE)
  gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)

  image = cut(image, gray)

  cv.imshow('output', image)
  cv.imwrite('output.png', image)
  while cv.waitKey(1) != ord('q'): pass


if __name__ == '__main__':
  main(sys.argv[1] if len(sys.argv) > 1 else 'lab.jpeg')