Segment an image using python and PIL to calculate centroid and rotations of multiple rectangular objects

Question

I am using python and PIL to find the centroid and rotation of various rectangles (and squares) in a 640x480 image, similar to this one

Answer 1

You need to identify each object before finding the corners. You only need the border of the objects, so you could also reduce your initial input to that. Then it is only a matter of following each distinct border to find your corners, the centroid is directly found after you know each distinct border.

Using the code below, here is what you get (centroid is the red point, the white text is the rotation in degrees):

Note that your input is not binary, so I used a really simple threshold for that. Also, the following code is the simplest way to achieve this, there are faster methods in any decent library.

import sys
import math
from PIL import Image, ImageOps, ImageDraw

orig = ImageOps.grayscale(Image.open(sys.argv[1]))
orig_bin = orig.point(lambda x: 0 if x < 128 else 255)
im = orig_bin.load()

border = Image.new('1', orig.size, 'white')
width, height = orig.size
bim = border.load()
# Keep only border points
for x in xrange(width):
    for y in xrange(height):
        if im[x, y] == 255:
            continue
        if im[x+1, y] or im[x-1, y] or im[x, y+1] or im[x, y-1]:
            bim[x, y] = 0
        else:
            bim[x, y] = 255

# Find each border (the trivial dummy way).
def follow_border(im, x, y, used):
    work = [(x, y)]
    border = []
    while work:
        x, y = work.pop()
        used.add((x, y))
        border.append((x, y))
        for dx, dy in ((1, 0), (-1, 0), (0, 1), (0, -1),
                (1, 1), (-1, -1), (1, -1), (-1, 1)):
            px, py = x + dx, y + dy
            if im[px, py] == 255 or (px, py) in used:
                continue
            work.append((px, py))

    return border

used = set()
border = []
for x in xrange(width):
    for y in xrange(height):
        if bim[x, y] == 255 or (x, y) in used:
            continue
        b = follow_border(bim, x, y, used)
        border.append(b)

# Find the corners and centroid of each rectangle.
rectangle = []
for b in border:
    xmin, xmax, ymin, ymax = width, 0, height, 0
    mean_x, mean_y = 0, 0
    b = sorted(b)
    top_left, bottom_right = b[0], b[-1]
    for x, y in b:
        mean_x += x
        mean_y += y
    centroid = (mean_x / float(len(b)), mean_y / float(len(b)))
    b = sorted(b, key=lambda x: x[1])
    curr = 0
    while b[curr][1] == b[curr + 1][1]:
        curr += 1
    top_right = b[curr]
    curr = len(b) - 1
    while b[curr][1] == b[curr - 1][1]:
        curr -= 1
    bottom_left = b[curr]

    rectangle.append([
        [top_left, top_right, bottom_right, bottom_left], centroid])


result = orig.convert('RGB')
draw = ImageDraw.Draw(result)
for corner, centroid in rectangle:
    draw.line(corner + [corner[0]], fill='red', width=2)
    cx, cy = centroid
    draw.ellipse((cx - 2, cy - 2, cx + 2, cy + 2), fill='red')
    rotation = math.atan2(corner[0][1] - corner[1][1],
            corner[1][0] - corner[0][0])
    rdeg = math.degrees(rotation)
    draw.text((cx + 10, cy), text='%.2f' % rdeg)

result.save(sys.argv[2])