facialRecognitionDetection2D.py

'''
  * ************************************************************
  *      Program: Facial Recognition Detection 2D
  *      Type: Python
  *      Author: David Velasco Garcia @davidvelascogarcia
  * ************************************************************
  *
  * | INPUT PORT                           | CONTENT                                                 |
  * |--------------------------------------|---------------------------------------------------------|
  * | /facialRecognitionDetection2D/img:i  | Input image                                             |
  *
  *
  * | OUTPUT PORT                          | CONTENT                                                 |
  * |--------------------------------------|---------------------------------------------------------|
  * | /facialRecognitionDetection2D/img:o  | Output image with facial detection                      |
  * | /facialRecognitionDetection2D/data:o | Output result, facial recognition data                  |
  * | /facialRecognitionDetection2D/coord:o| Output result, facial recognition coordinates           |
'''

# Libraries
import configparser
import cv2
import datetime
import face_recognition
from halo import Halo
import numpy as np
import platform
import queue
import threading
import time
import yarp


class FacialRecognitionDetection2D:

    # Function: Constructor
    def __init__(self):

        # Build Halo spinner
        self.systemResponse = Halo(spinner='dots')

    # Function: getSystemPlatform
    def getSystemPlatform(self):

        # Get system configuration
        print("\nDetecting system and release version ...\n")
        systemPlatform = platform.system()
        systemRelease = platform.release()

        print("**************************************************************************")
        print("Configuration detected:")
        print("**************************************************************************")
        print("\nPlatform:")
        print(systemPlatform)
        print("Release:")
        print(systemRelease)

        return systemPlatform, systemRelease

    # Function: getAuthenticationData
    def getAuthenticationData(self):

        print("\n**************************************************************************")
        print("Authentication:")
        print("**************************************************************************\n")

        loopControlFileExists = 0

        while int(loopControlFileExists) == 0:
            try:
                # Get authentication data
                print("\nGetting authentication data ...\n")

                authenticationData = configparser.ConfigParser()
                authenticationData.read('../config/config.ini')
                authenticationData.sections()

                imageWidth = authenticationData['Configuration']['image-width']
                imageHeight = authenticationData['Configuration']['image-height']

                print("Image width: " + str(imageWidth))
                print("Image height: " + str(imageHeight))

                # Exit loop
                loopControlFileExists = 1

            except:

                systemResponseMessage = "\n[ERROR] Sorry, config.ini not founded, waiting 4 seconds to the next check ...\n"
                self.systemResponse.text_color = "red"
                self.systemResponse.fail(systemResponseMessage)
                time.sleep(4)

        systemResponseMessage = "\n[INFO] Data obtained correctly.\n"
        self.systemResponse.text_color = "green"
        self.systemResponse.succeed(systemResponseMessage)

        return imageWidth, imageHeight

    # Function: getImageDatabase
    def getImageDatabase(self):

        loopControlFileExists = 0

        while int(loopControlFileExists) == 0:
            try:
                # Get authentication data
                print("\nGetting image database ...\n")

                # Read image database from file
                imageDatabaseFile = open('./../resources/imageDatabase.txt', 'r')
                imageDatabaseLines = imageDatabaseFile.readlines()

                # Prepare image array and variable to count number of images
                imageNumber = 0
                imageDatabase = []

                # Append images files to imageDatabase array
                for imageDatabaseLine in imageDatabaseLines:

                    print("Image " + str(imageNumber) + ": " + str(imageDatabaseLine.strip()))

                    # Prepare image path
                    imagePath = "./../database/" + imageDatabaseLine.strip()

                    # Append image path
                    imageDatabase.append(imagePath)

                    # Increase image number
                    imageNumber = imageNumber + 1

                # Exit loop
                loopControlFileExists = 1

            except:

                systemResponseMessage = "\n[ERROR] Sorry, imageDatabase.txt not founded, waiting 4 seconds to the next check ...\n"
                self.systemResponse.text_color = "red"
                self.systemResponse.fail(systemResponseMessage)
                time.sleep(4)

        systemResponseMessage = "\n[INFO] Data obtained correctly.\n"
        self.systemResponse.text_color = "green"
        self.systemResponse.succeed(systemResponseMessage)

        return imageDatabase

    # Function: getNameDatabase
    def getNameDatabase(self):

        loopControlFileExists = 0

        while int(loopControlFileExists) == 0:
            try:
                # Get authentication data
                print("\nGetting users name ...\n")

                # Read name database from file
                nameDatabaseFile = open('./../resources/nameDatabase.txt', 'r')
                nameDatabaseLines = nameDatabaseFile.readlines()

                # Prepare name array and variable to count number of names
                nameNumber = 0
                nameDatabase = []

                # Append names files to nameDatabase array
                for nameDatabaseLine in nameDatabaseLines:
                    print("Name " + str(nameNumber) + ": " + str(nameDatabaseLine.strip()))

                    # Prepare name path
                    namePath = nameDatabaseLine.strip()

                    # Append name path
                    nameDatabase.append(namePath)

                    # Increase name number
                    nameNumber = nameNumber + 1

                # Exit loop
                loopControlFileExists = 1

            except:

                systemResponseMessage = "\n[ERROR] Sorry, nameDatabase.txt not founded, waiting 4 seconds to the next check ...\n"
                self.systemResponse.text_color = "red"
                self.systemResponse.fail(systemResponseMessage)
                time.sleep(4)

        systemResponseMessage = "\n[INFO] Data obtained correctly.\n"
        self.systemResponse.text_color = "green"
        self.systemResponse.succeed(systemResponseMessage)

        return nameDatabase, nameNumber

    # Function: trainModel
    def trainModel(self, imageDatabase, nameDatabase, nameNumber):

        userDatabase = []
        userDatabaseEncoded = []
        userTrainedDatabase = []
        nameTrainedDatabase = []

        # Prepare image database and encoded image database
        for user in range(int(nameNumber)):

            # Set user database images form image database
            userDatabase.append(face_recognition.load_image_file(imageDatabase[user]))

            # Detect user face in user database images and encode in user database encoded with training model
            userDatabaseEncoded.append(face_recognition.face_encodings(userDatabase[user])[0])

            # Prepare database with trained users and names
            userTrainedDatabase.append(userDatabaseEncoded[user])
            nameTrainedDatabase.append(nameDatabase[user])

        return userTrainedDatabase, nameTrainedDatabase

    # Function: getFacesCompare
    def getFacesCompare(self, userTrainedDatabase, faceEncoding, compareFacesQueueBuffer):

        # Compare face detected with user trained database
        match = face_recognition.compare_faces(userTrainedDatabase, faceEncoding)

        # Put in Queue buffer
        compareFacesQueueBuffer.put(match)

    # Function: compareFaces
    def getFacesDistance(self, userTrainedDatabase, faceEncoding, distanceFacesQueueBuffer):

        # Compare distance between face detected and trained database
        matchDistance = face_recognition.face_distance(userTrainedDatabase, faceEncoding)

        # Put in Queue buffer
        distanceFacesQueueBuffer.put(matchDistance)

    # Function: analyzeImage
    def analyzeImage(self, dataToSolve, userTrainedDatabase, nameTrainedDatabase, imageWidth, imageHeight, inputImagePort, outputImagePort, outputDataPort):

        # Detect faces in data to solve and get coordinates of all them
        faceLocations = face_recognition.face_locations(dataToSolve)

        # Get detected faces of coordinates
        faceEncodings = face_recognition.face_encodings(dataToSolve, faceLocations)

        # If a face is detected
        if str(faceLocations) != "[]":

            # Compare each face detected in data to solve with user trained database
            for (yMax, xMax, yMin, xMin), faceEncoding in zip(faceLocations, faceEncodings):

                # Create Queue buffers
                compareFacesQueueBuffer = queue.Queue()
                distanceFacesQueueBuffer = queue.Queue()

                # Create threads
                compareFacesThread = threading.Thread(target=self.getFacesCompare, args=(userTrainedDatabase, faceEncoding, compareFacesQueueBuffer))
                distanceFacesThread = threading.Thread(target=self.getFacesDistance, args=(userTrainedDatabase, faceEncoding, distanceFacesQueueBuffer))

                # Start thread
                compareFacesThread.start()
                distanceFacesThread.start()

                # Wait until both threads ends
                compareFacesThread.join()
                distanceFacesThread.join()

                # Get results of threads
                match = compareFacesQueueBuffer.get()
                matchDistance = distanceFacesQueueBuffer.get()

                # Extract the most certain index based on match distance
                predictionIndex = np.argmin(matchDistance)

                # If detected user is in database get the name
                if match[predictionIndex]:
                    detectedUser = nameTrainedDatabase[predictionIndex]

                # If isn´t in database is "Unknown"
                else:
                    detectedUser = "Unknown"

                # Draw a blue rectangle in detected face
                cv2.rectangle(dataToSolve, (xMin, yMax), (xMax, yMin), (0, 0, 255), 2)

                # Draw a blue rectangle at the bottom of the main rectangle that will contain detected user
                cv2.rectangle(dataToSolve, (xMin, yMin - 35), (xMax, yMin), (0, 0, 255), cv2.FILLED)

                # Draw the user detected in previously blue rectangle
                cv2.putText(dataToSolve, detectedUser, (xMin + 6, yMin - 6), cv2.FONT_HERSHEY_DUPLEX, 1.0, (255, 255, 255), 1)

                # Get coordinates
                coordinatesXY = self.getCoordinates(xMin, yMin, imageHeight)

                # Prepare output results
                dataSolvedResults = "Detection: " + str(detectedUser) + ", Coordinates: " + str(coordinatesXY) + ", Date: " + str(datetime.datetime.now())

                # Send detected results
                outputDataPort.send(dataSolvedResults)

        else:
            systemResponseMessage = "\n[INFO] No faces detected.\n"
            self.systemResponse.text_color = "blue"
            self.systemResponse.info(systemResponseMessage)

            # Prepare output results
            dataSolvedResults = "Detection: None, Coordinates: None, Date: " + str(datetime.datetime.now())

            # Send detected results
            outputDataPort.send(dataSolvedResults)

        systemResponseMessage = "\n" + str(dataSolvedResults) + "\n"
        self.systemResponse.text_color = "green"
        self.systemResponse.succeed(systemResponseMessage)

        dataSolvedImage = dataToSolve

        return dataSolvedImage

    # Function: getCoordinates
    def getCoordinates(self, xMin, yMin, imageHeight):

        # Get centroid coordinates
        x = xMin
        y = int(imageHeight) - yMin

        # Prepare coordinates
        coordinatesXY = str(x) + ", " + str(y)

        return coordinatesXY

    # Function: processRequest
    def processRequests(self, userTrainedDatabase, nameTrainedDatabase, imageWidth, imageHeight, inputImagePort, outputImagePort, outputDataPort):

        # Variable to control loopProcessRequests
        loopProcessRequests = 0

        while int(loopProcessRequests) == 0:

            # Waiting to input data request
            print("**************************************************************************")
            print("Waiting for input data request:")
            print("**************************************************************************")

            systemResponseMessage = "\n[INFO] Waiting for input data request at " + str(datetime.datetime.now()) + " ...\n"
            self.systemResponse.text_color = "yellow"
            self.systemResponse.warn(systemResponseMessage)

            # Receive input request
            dataToSolve = inputImagePort.receive()

            print("\n**************************************************************************")
            print("Processing:")
            print("**************************************************************************\n")

            try:
                dataSolvedImage = self.analyzeImage(dataToSolve, userTrainedDatabase, nameTrainedDatabase, imageWidth, imageHeight, inputImagePort, outputImagePort, outputDataPort)

                # Send output results
                outputImagePort.send(dataSolvedImage)

            except:
                systemResponseMessage = "\n[ERROR] Sorry, i couldn´t resolve your request.\n"
                self.systemResponse.text_color = "red"
                self.systemResponse.fail(systemResponseMessage)


class YarpDataPort:

    # Function: Constructor
    def __init__(self, portName):

        # Build Halo spinner
        self.systemResponse = Halo(spinner='dots')

        # Build port and bottle
        self.yarpPort = yarp.Port()
        self.yarpBottle = yarp.Bottle()

        systemResponseMessage = "\n[INFO] Opening Yarp data port " + str(portName) + " ...\n"
        self.systemResponse.text_color = "yellow"
        self.systemResponse.warn(systemResponseMessage)

        # Open Yarp port
        self.portName = portName
        self.yarpPort.open(self.portName)

    # Function: receive
    def receive(self):

        self.yarpPort.read(self.yarpBottle)
        dataReceived = self.yarpBottle.toString()
        dataReceived = dataReceived.replace('"', '')

        systemResponseMessage = "\n[RECEIVED] Data received: " + str(dataReceived) + " at " + str(datetime.datetime.now()) + ".\n"
        self.systemResponse.text_color = "blue"
        self.systemResponse.info(systemResponseMessage)

        return dataReceived

    # Function: send
    def send(self, dataToSend):

        self.yarpBottle.clear()
        self.yarpBottle.addString(str(dataToSend))
        self.yarpPort.write(self.yarpBottle)

    # Function: close
    def close(self):

        systemResponseMessage = "\n[INFO] " + str(self.portName) + " port closed correctly.\n"
        self.systemResponse.text_color = "yellow"
        self.systemResponse.warn(systemResponseMessage)

        self.yarpPort.close()


class YarpImagePort:

    # Function: Constructor
    def __init__(self, portName, imageWidth, imageHeight):

        # Build Halo spinner
        self.systemResponse = Halo(spinner='dots')

        # If input image port required
        if "/img:i" in str(portName):
            self.yarpPort = yarp.BufferedPortImageRgb()

        # If output image port required
        else:
            self.yarpPort = yarp.Port()

        systemResponseMessage = "\n[INFO] Opening Yarp image port " + str(portName) + " ...\n"
        self.systemResponse.text_color = "yellow"
        self.systemResponse.warn(systemResponseMessage)

        # Open Yarp port
        self.portName = portName
        self.yarpPort.open(self.portName)

        # Build image buffer
        self.imageWidth = int(imageWidth)
        self.imageHeight = int(imageHeight)
        self.bufferImage = yarp.ImageRgb()
        self.bufferImage.resize(self.imageWidth, self.imageHeight)
        self.bufferArray = np.ones((self.imageHeight, self.imageWidth, 3), np.uint8)
        self.bufferImage.setExternal(self.bufferArray.data, self.bufferArray.shape[1], self.bufferArray.shape[0])

    # Function: receive
    def receive(self):

        image = self.yarpPort.read()
        self.bufferImage.copy(image)
        assert self.bufferArray.__array_interface__['data'][0] == self.bufferImage.getRawImage().__int__()
        image = self.bufferArray[:, :, ::-1]

        return self.bufferArray

    # Function: send
    def send(self, dataToSend):

        self.bufferArray[:,:] = dataToSend
        self.yarpPort.write(self.bufferImage)

    # Function: close
    def close(self):

        systemResponseMessage = "\n[INFO] " + str(self.portName) + " port closed correctly.\n"
        self.systemResponse.text_color = "yellow"
        self.systemResponse.warn(systemResponseMessage)

        self.yarpPort.close()


# Function: main
def main():

    print("**************************************************************************")
    print("**************************************************************************")
    print("              Program: Facial Recognition Detection 2D                    ")
    print("                     Author: David Velasco Garcia                         ")
    print("                             @davidvelascogarcia                          ")
    print("**************************************************************************")
    print("**************************************************************************")

    print("\nLoading Facial Recognition Detection 2D engine ...\n")

    # Build facialRecognitionDetection2D object
    facialRecognitionDetection2D = FacialRecognitionDetection2D()

    # Get system platform
    systemPlatform, systemRelease = facialRecognitionDetection2D.getSystemPlatform()

    # Get authentication data
    imageWidth, imageHeight = facialRecognitionDetection2D.getAuthenticationData()

    # Get image database
    imageDatabase = facialRecognitionDetection2D.getImageDatabase()

    # Get name database
    nameDatabase, nameNuber = facialRecognitionDetection2D.getNameDatabase()

    # Train model
    userTrainedDatabase, nameTrainedDatabase = facialRecognitionDetection2D.trainModel(imageDatabase, nameDatabase, nameNuber)

    # Init Yarp network
    yarp.Network.init()

    # Create Yarp ports
    inputImagePort = YarpImagePort("/facialRecognitionDetection2D/img:i", imageWidth, imageHeight)
    outputImagePort = YarpImagePort("/facialRecognitionDetection2D/img:o", imageWidth, imageHeight)
    outputDataPort = YarpDataPort("/facialRecognitionDetection2D/data:o")

    # Process input requests
    facialRecognitionDetection2D.processRequests(userTrainedDatabase, nameTrainedDatabase, imageWidth, imageHeight, inputImagePort, outputImagePort, outputDataPort)

    # Close Yarp ports
    inputImagePort.close()
    outputImagePort.close()
    outputDataPort.close()

    print("**************************************************************************")
    print("Program finished")
    print("**************************************************************************")
    print("\nfacialRecognitionDetection2D program finished correctly.\n")


if __name__ == "__main__":

    # Call main function
    main()