Augmented Reality using OpenCV and Python

Tags

, , , , , , , , , , , , ,

Arkwood has a long-standing phobia of Lego policemen. ‘They are watching me!’ he shrieked, the hairs on the back of his neck erect. Of course, it was just his drug-addled brain playing a cruel trick.

‘Anyway,’ I said, ‘it’s the Lego criminals you need to watch out for. Their sticky yellow fingers are oft found in the cookie jar.’

In my last post, OpenCV Camera Calibration and Pose Estimation using Python, I was able to calibrate my webcam and draw a 3D cube on a grid. And in post Lego detection using OpenCV (Mark III) I was able to detect Lego policemen using a OpenCV haar cascade classifier. Let’s bring the two together and provide some augmented reality!

Here’s the code:

from webcam import Webcam
from detection import Detection
from effects import Effects
import cv2

# set up classes
webcam = Webcam()
webcam.start()
detection = Detection()
effects = Effects()
 
# loop forever
while True:
  
    # if lego policeman detected...
    image = webcam.get_current_frame()
    item_detected = detection.is_item_detected_in_image('haarcascade_lego_policeman.xml', image)
  
    # ...then draw a virtual jail for all those nasty criminals
    if item_detected:
        effects.render(image)
                     
    # show the scene
    cv2.imshow('scene',image)
    cv2.waitKey(100)

Dropping into a while loop, I obtain the current frame from my webcam via the Webcam class.

Next, the Detection class tries to find Lego policemen in the webcam image, using a haar cascade classifier.

If a Lego policeman has been found then I use the Effects class to draw a 3D jail on the image, to incarcerate any Lego criminals.

Finally, we show the augmented image in a window.

The code from the classes, along with a link to my haar cascade classifier, is at the foot of this post. But first, a demo…

The scene is set:

augmentedreality_scene

A Lego soldier walks into shot. A military man is not an officer of the law – alas – therefore our grid does not yield a virtual jail:

augmentedreality_soldier

Who’s that strolling up the boulevard? It’s a nurse. Sorry, my angel of the wards, but the virtual jail is not for you:

augmentedreality_nurse

But what is that I hear? Why, it is a bobby on the beat:

augmentedreality_policeman

Hurray! Our haar cascade classifier has detected the Lego policeman in the webcam image. A virtual jail has been drawn on the grid.

We can build and pan our jail at different angles:

augmentedreality_policeman_2

augmentedreality_policeman_3

Granted, there is a bit of work still to do. For a start, the jail has no bars, so all those nasty criminals will be able to make good their escape. The grid I’ve used to render the jail is a bit big, taking up most of the scene. And we could do with omitting the detection rectangle around the policeman, as it spoils the sense of mind-boggling augmented reality.

Still, it’s a start.

‘Don’t worry Arkwood,’ I sympathised, ‘I’ll soon have a proper jail for the Lego policemen. You won’t have to worry about them watching you, as they’ll be too busy collaring all those crooks.’

What Madness!

P.S.

Here’s the code I promised. First, the Webcam class, which runs in a thread:

import cv2
from threading import Thread
 
class Webcam:
 
    def __init__(self):
        self.video_capture = cv2.VideoCapture(0)
        self.current_frame = self.video_capture.read()[1]
         
    # create thread for capturing images
    def start(self):
        Thread(target=self._update_frame, args=()).start()
 
    def _update_frame(self):
        while(True):
            self.current_frame = self.video_capture.read()[1]
                 
    # get the current frame
    def get_current_frame(self):
        return self.current_frame

Next, the Detection class. It uses my Lego Policeman haar cascade to attempt to find Lego policemen in the supplied image.

import cv2

class Detection(object):

    def is_item_detected_in_image(self, item_cascade_path, image):
 
        # do detection
        item_cascade = cv2.CascadeClassifier(item_cascade_path)
        gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        items = item_cascade.detectMultiScale(gray_image, scaleFactor=1.1, minNeighbors=36)

        for (x,y,w,h) in items:
            cv2.rectangle(image,(x,y),(x+w,y+h),(255,0,0),2)
               
        # indicate whether item detected in image
        return len(items) > 0

And finally, my Effects class, which borrows heavily from the OpenCV Pose Estimation article:

import cv2
import numpy as np
 
class Effects(object):
   
    def render(self, image):
 
        # load calibration data
        with np.load('webcam_calibration_ouput.npz') as X:
            mtx, dist, _, _ = [X[i] for i in ('mtx','dist','rvecs','tvecs')]
 
        # set up criteria, object points and axis
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
         
        objp = np.zeros((6*7,3), np.float32)
        objp[:,:2] = np.mgrid[0:7,0:6].T.reshape(-1,2)
 
        axis = np.float32([[0,0,0], [0,3,0], [3,3,0], [3,0,0],
                           [0,0,-3],[0,3,-3],[3,3,-3],[3,0,-3] ])
 
        # find grid corners in image
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        ret, corners = cv2.findChessboardCorners(gray, (7,6), None)
 
        if ret == True:
             
            # project 3D points to image plane
            cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
            rvecs, tvecs, inliers = cv2.solvePnPRansac(objp, corners, mtx, dist)
            imgpts, jac = cv2.projectPoints(axis, rvecs, tvecs, mtx, dist)
 
            # draw cube
            self._draw_cube(image, corners, imgpts)
 
    def _draw_cube(self, img, corners, imgpts):
        imgpts = np.int32(imgpts).reshape(-1,2)
 
        # draw floor
        cv2.drawContours(img, [imgpts[:4]],-1,(200,150,10),-3)
 
        # draw pillars
        for i,j in zip(range(4),range(4,8)):
            cv2.line(img, tuple(imgpts[i]), tuple(imgpts[j]),(255),3)
 
        # draw roof
        cv2.drawContours(img, [imgpts[4:]],-1,(200,150,10),3)
Follow

Get every new post delivered to your Inbox.

Join 80 other followers