Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Updated for Python 3 #3

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Updated for Python 3 #3

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Updated for Python 3 
Small updates so this runs in Python 3. Mainly updating print statements (including an import to retain backwards compatibility) and fixing float conversions to ints for the box drawing. Only lightly tested!
  • Loading branch information
@robertlayton
robertlayton authored Jun 23, 2020
commit a5c64144ebdace9b089f1d2d595f81bff2ba3c5e
38 changes: 31 additions & 7 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from __future__ import print_function
from PIL import Image, ImageDraw
from collections import Counter
import heapq
Expand All @@ -17,20 +18,23 @@
AREA_POWER = 0.25
OUTPUT_SCALE = 1


def weighted_average(hist):
total = sum(hist)
value = sum(i * x for i, x in enumerate(hist)) / total
error = sum(x * (value - i) ** 2 for i, x in enumerate(hist)) / total
error = error ** 0.5
return value, error


def color_from_histogram(hist):
r, re = weighted_average(hist[:256])
g, ge = weighted_average(hist[256:512])
b, be = weighted_average(hist[512:768])
e = re * 0.2989 + ge * 0.5870 + be * 0.1140
return (r, g, b), e


def rounded_rectangle(draw, box, radius, color):
l, t, r, b = box
d = radius * 2
Expand All @@ -42,6 +46,7 @@ def rounded_rectangle(draw, box, radius, color):
draw.rectangle((l, t + d, r, b - d), color)
draw.rectangle((l + d, t, r - d, b), color)


class Quad(object):
def __init__(self, model, box, depth):
self.model = model
Expand All @@ -52,12 +57,15 @@ def __init__(self, model, box, depth):
self.leaf = self.is_leaf()
self.area = self.compute_area()
self.children = []

def is_leaf(self):
l, t, r, b = self.box
return int(r - l <= LEAF_SIZE or b - t <= LEAF_SIZE)

def compute_area(self):
l, t, r, b = self.box
return (r - l) * (b - t)

def split(self):
l, t, r, b = self.box
lr = l + (r - l) / 2
Expand All @@ -69,6 +77,7 @@ def split(self):
br = Quad(self.model, (lr, tb, r, b), depth)
self.children = (tl, tr, bl, br)
return self.children

def get_leaf_nodes(self, max_depth=None):
if not self.children:
return [self]
Expand All @@ -79,6 +88,7 @@ def get_leaf_nodes(self, max_depth=None):
result.extend(child.get_leaf_nodes(max_depth))
return result


class Model(object):
def __init__(self, path):
self.im = Image.open(path).convert('RGB')
Expand All @@ -87,23 +97,29 @@ def __init__(self, path):
self.root = Quad(self, (0, 0, self.width, self.height), 0)
self.error_sum = self.root.error * self.root.area
self.push(self.root)

@property
def quads(self):
return [x[-1] for x in self.heap]

def average_error(self):
return self.error_sum / (self.width * self.height)

def push(self, quad):
score = -quad.error * (quad.area ** AREA_POWER)
heapq.heappush(self.heap, (quad.leaf, score, quad))

def pop(self):
return heapq.heappop(self.heap)[-1]

def split(self):
quad = self.pop()
self.error_sum -= quad.error * quad.area
children = quad.split()
for child in children:
self.push(child)
self.error_sum += child.error * child.area

def render(self, path, max_depth=None):
m = OUTPUT_SCALE
dx, dy = (PADDING, PADDING)
Expand All @@ -113,43 +129,51 @@ def render(self, path, max_depth=None):
for quad in self.root.get_leaf_nodes(max_depth):
l, t, r, b = quad.box
box = (l * m + dx, t * m + dy, r * m - 1, b * m - 1)

# Convert box and color to ints. Python2 did this automatically
box = tuple(int(v) for v in box)
quad.color = tuple(int(v) for v in quad.color)

if MODE == MODE_ELLIPSE:
draw.ellipse(box, quad.color)
elif MODE == MODE_ROUNDED_RECTANGLE:
radius = m * min((r - l), (b - t)) / 4
# Compute radius and convert to ints
radius = int(m * min((r - l), (b - t)) / 4)
rounded_rectangle(draw, box, radius, quad.color)
else:
draw.rectangle(box, quad.color)
del draw
im.save(path, 'PNG')


def main():
args = sys.argv[1:]
if len(args) != 1:
print 'Usage: python main.py input_image'
print('Usage: python main.py input_image')
return
model = Model(args[0])
previous = None
for i in range(ITERATIONS):
error = model.average_error()
if previous is None or previous - error > ERROR_RATE:
print i, error
print(i, error)
if SAVE_FRAMES:
model.render('frames/%06d.png' % i)
previous = error
model.split()
model.render('output.png')
print '-' * 32
print('-' * 32)
depth = Counter(x.depth for x in model.quads)
for key in sorted(depth):
value = depth[key]
n = 4 ** key
pct = 100.0 * value / n
print '%3d %8d %8d %8.2f%%' % (key, n, value, pct)
print '-' * 32
print ' %8d %8.2f%%' % (len(model.quads), 100)
print('%3d %8d %8d %8.2f%%' % (key, n, value, pct))
print('-' * 32)
print(' %8d %8.2f%%' % (len(model.quads), 100))
# for max_depth in range(max(depth.keys()) + 1):
# model.render('out%d.png' % max_depth, max_depth)


if __name__ == '__main__':
main()