6
\$\begingroup\$

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:

enter image description here

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.

My code is hosted on github here: AR-Testing

You can download all the code with the test images here: AR-Testing.zip

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.