filter.py

filter.py

#!/usr/bin/env python                                                                                                   

import os
import pickle
import random
import sys

import beanstalkc

# Use *'s sparingly!
from opencv.cv import *
from opencv.highgui import *

TEMP_DIRECTORY = '/tmp'

MQHOST = 'arbiter'
MQPORT = 11300

def main():
    # Read from newImages, write to people
    beanstalk = beanstalkc.Connection(host=MQHOST, port=MQPORT)
    beanstalk.use('people')
    beanstalk.watch('newImages')

    while True:
        # Request a job from the arbiter
        job = beanstalk.reserve()

        # Deserialize the image and write to a temp file
        filename, data = pickle.loads(job.body)
        image = open(os.path.join(TEMP_DIRECTORY, filename), 'w')
        image.write(data)
        image.close()

        # Load the image with OpenCV
        image = cvLoadImage(os.path.join(TEMP_DIRECTORY, filename))
        if image:
            # Convert to grayscale for face detection
            grayscale = cvCreateImage(cvSize(image.width, image.height), 8, 1)
            cvCvtColor(image, grayscale, CV_BGR2GRAY)

            storage = cvCreateMemStorage(0)
            cvClearMemStorage(storage)
            cvEqualizeHist(grayscale, grayscale)

            # Pass grayscale image through face detection filter
            cascade = cvLoadHaarClassifierCascade(
                '/usr/share/opencv/haarcascades/haarcascade_frontalface_default.xml',
                cvSize(1,1))
            faces = cvHaarDetectObjects(grayscale, cascade, storage, 1.2, 2,
                                        CV_HAAR_DO_CANNY_PRUNING, cvSize(50,50))

            draw = ""
            if faces.total > 0:
                # Build arguments for imagemagick
                for f in faces:
                    draw += 'rectangle %d,%d %d,%d ' %(f.x, f.y, f.x+f.width, f.y+f.height)

                    # Draw rectangles over each detected face
                    os.system('convert %s -stroke red -fill none -draw "%s" %s' % (
                            os.path.join(TEMP_DIRECTORY, filename), draw,
                            os.path.join(TEMP_DIRECTORY, "h." + filename)))

                    with open(os.path.join(TEMP_DIRECTORY, "h." + filename)) as data_h:
                        data = data_h.read()
                        beanstalk.put(pickle.dumps((filename, data)))

                # Delete temporary files
                os.unlink(os.path.join(TEMP_DIRECTORY, "h." + filename))
            os.unlink(os.path.join(TEMP_DIRECTORY, filename))
        # Tell the arbiter that the job was successfully completed
        job.delete()

if __name__ == '__main__':
    main()