I have been working on creating a simple AR Tag classifier, for detecting a simple AR Tag glyph. Spitting out an image that looks like this:
With the red circle indicating the corner closest to the orientation square.
My program is structured as follows:
AR.py (where all the magic happens)
import cv2
import numpy as np
from artag import ARtag
def display_scaled_image(name, image, scale):
"""Function to display a scaled cv2 image
:param name:
Window name
:type name:
basestring
:param image:
Image as numpy array
:type image:
numpy.ndarray
:param scale:
Scale factor applied to image
:type scale:
float
"""
height, width = image.shape[:2]
cv2.imshow(name, cv2.resize(image,
(int(scale * width), int(scale * height)),
interpolation=cv2.INTER_CUBIC))
def mask_black(img):
"""Function to mask for black boxes in an BGR image
:param img:
Numpy image to process
:type img:
numpy.ndarray
:return:
the masked image
:rtype:
numpy.ndarray
"""
return cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
cv2.THRESH_BINARY, 511, 10)
def find_tag(img):
"""
Identify AR tags in an image
:param img:
The image to be processed, may be pre-converted to greyscale
:type img:
numpy.ndarray
:return:
Tuple[List[artag], numpy.ndarray]
"""
# Convert the image to greyscale if it is not
if len(img.shape) > 2:
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
dims = img.shape[:2]
# Normalize image size
img, _ = rescale_img(img)
kernel = np.ones((3, 3), np.uint8)
img = mask_black(img)
img = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel)
_, contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
# Clean up contours to regular shapes
culled_conts = []
for c, h in zip(contours, hierarchy[0]):
c = cv2.convexHull(c)
epsilon = 0.01 * cv2.arcLength(c, True)
approx = cv2.approxPolyDP(c, epsilon, True)
culled_conts.append((approx, h))
img, scale = rescale_img(img, dims)
# Identify nested shapes as possible AR tags
ar_tags = []
for c, h in culled_conts:
# 4 sided tags
if len(c) == 4:
# The tag contains a smaller contour
if h[2] != -1:
tag = ARtag(c, culled_conts[h[2]][0], scale)
if tag.valid():
ar_tags.append(tag)
img = cv2.merge((img, img, img))
return ar_tags, img
def rescale_img(img, dims=None, width=1024):
"""Rescale image, either to the dimensions provided or to the width
provided, returning the image and it's dimensions
:param img:
The input image
:type img:
numpy.ndarray
:param dims:
Provide this parameter to fit the image to a specific size
:type dims:
Tuple[int, int]
:param width:
The width to resize the image to, (overridden by dims) deafualts to 1024
:type width:
int
:return:
Tuple of image and its dimensions
:rtype:
Tuple[numpy.ndarray, Tuple[int, int]]
"""
grey = img
if dims is None:
r = float(width) / grey.shape[1]
dims = (width, int(grey.shape[0] * r))
else:
dims = dims[::-1]
scale = (dims[1] / float(grey.shape[0]), dims[0] / float(grey.shape[1]))
return cv2.resize(grey, dims, interpolation=cv2.INTER_AREA), scale
if __name__ == "__main__":
# Display results for all test images
import os
for f in list(os.walk("Test imges"))[0][2]:
if "jpg" in f:
in_img = cv2.imread("Test imges/" + f)
grey = cv2.cvtColor(in_img, cv2.COLOR_BGR2GRAY)
ar_tags, proc_img = find_tag(grey)
for i in ar_tags:
proc_img = i.draw(proc_img)
in_img = i.draw(in_img)
display_scaled_image("Image - {}".format(f), proc_img, .25)
display_scaled_image("Orig Image - {}".format(f), in_img, .25)
cv2.waitKey(0)
cv2.destroyAllWindows()
artag.py
import math
import cv2
class ARtag(object):
def __init__(self, outer, inner, scale):
"""Create a ARTag object, based on its inner contour,
outer contour and the scale of the image processed as
related to the image that was processed to extract them
:param outer:
The outer contour (Assumed to have 4 sides)
:param inner:
The outer contour (Assumed to have at least 4 sides)
:param scale:
The scale of the image that was processed to create this tag
:type scale:
Tuple[int, int]
"""
self.outer_cont = outer
for i, c in enumerate(self.outer_cont):
self.outer_cont[i][0][0] = c[0][0] * scale[1]
self.outer_cont[i][0][1] = c[0][1] * scale[0]
self.inner_cont = inner
for i, c in enumerate(self.inner_cont):
self.inner_cont[i][0][0] = c[0][0] * scale[1]
self.inner_cont[i][0][1] = c[0][1] * scale[0]
self.outer_tup = [tuple(i[0]) for i in outer]
self.inner_tup = [tuple(i[0]) for i in inner]
x, y = zip(*self.inner_tup)
x = sum(x) / len(x)
y = sum(y) / len(y)
self.outer_tup = sorted(self.outer_tup, key=lambda p: math.sqrt((p[0] - x)**2 + (p[1] - y)**2))
# Points of interest labeled clockwise from marked corner
self.p0 = self.outer_tup[0]
self.p1 = self.outer_tup[1]
self.p2 = self.outer_tup[3]
# Furthest point must be the 3rd point
self.p3 = self.outer_tup[2]
def draw(self, img):
"""Draw the tag on a 3 channel image
:param img:
The 3 channel image
:type img:
numpy.ndarray
:return:
The input image
:rtype:
numpy.ndarray
"""
img = cv2.drawContours(img, [self.outer_cont], -1, (0, 255, 0), 10)
img = cv2.drawContours(img, [self.inner_cont], -1, (255, 0, 0), 10)
img = cv2.circle(img, tuple(self.p0), 20, (0, 0, 255), 10)
return img
def ratio(self, img):
"""Ratio of the size of the Tag vs the input image size
:param img:
The input image
:type img:
numpy.ndarray
:return:
The processed image
:rtype:
numpy.ndarray
"""
return cv2.contourArea(self.outer_cont) / float(img.shape[0] * img.shape[1])
def in_outer(self):
return cv2.contourArea(self.inner_cont) / cv2.contourArea(self.outer_cont)
@staticmethod
def dist(p0, p1):
""" Return the distance between two tuples in the format (x, y)
:param p0:
The first point
:type p0:
Tuple[int, int]
:param p1:
The second point
:type p1:
Tuple[int, int]
:return:
The distance
:rtype:
float
"""
return math.sqrt((p0[0] - p1[0])**2 + (p0[1] - p1[1])**2)
def valid(self):
"""Check the validity of this artag
Check the following items:
- Any line distance is more than 1.8 times the average length
- Any line distance is less than .2 times the average length
- The inner contour takes up more than 10% of the tag
- The inner contour takes up less than 1% of the tag
:return:
The validity of the tag
:rtype:
bool
"""
lengths = [
self.dist(self.p0, self.p1),
self.dist(self.p1, self.p2),
self.dist(self.p2, self.p3),
self.dist(self.p3, self.p0)
]
# check aspect ratio
if any([i > 1.8 * sum(lengths) / len(lengths) for i in lengths]):
return False
if any([i < .2 * sum(lengths) / len(lengths) for i in lengths]):
return False
# Check area ratios
if cv2.contourArea(self.inner_cont) / cv2.contourArea(self.outer_cont) > 0.1:
return False
if cv2.contourArea(self.inner_cont) / cv2.contourArea(self.outer_cont) < 0.01:
return False
return True
It seems to work great, except for image 11 seems to fail to detect the Tag, this has somthing to do with the nesting of the contours. What I would like to do to fix this image, would be to do somthing like find large 4-sided contours that have smaller 4 sided contours inside, that are not nested and have 4 sides.
The problem is the structure of the hierarchy object returned by opencv. It has the indexes of the contours, but if you modify the list it makes this useless. I would love to have a data structure for the contours, that have the data baked in, but the problem occurs if I start removing contours, becouse they are invalid or somthing I get problems referring to contours that are nonexistent.
Any input on solving this problem elegantly would be a great in addition to a general review of the code.
