tensorflowLiteDetection2D.py

'''
  * ************************************************************
  *      Program: Tensorflow Lite Detection 2D Module
  *      Type: Python
  *      Author: David Velasco Garcia @davidvelascogarcia
  * ************************************************************
  *
  * | INPUT PORT                           | CONTENT                                                 |
  * |--------------------------------------|---------------------------------------------------------|
  * | /tensorflowLiteDetection2D/img:i     | Input image                                             |
  *
  *
  * | OUTPUT PORT                          | CONTENT                                                 |
  * |--------------------------------------|---------------------------------------------------------|
  * | /tensorflowLiteDetection2D/img:o     | Output image with detection                             |
  * | /tensorflowLiteDetection2D/data:o    | Output result, recognition data                         |
'''

# Libraries
import configparser
import cv2
import datetime
from halo import Halo
import importlib.util
import numpy as np
import platform
import time
import yarp

# Check tensorflow lite libraries
checkTensorflowLiteLibs = importlib.util.find_spec('tflite_runtime')

if checkTensorflowLiteLibs:
    from tflite_runtime.interpreter import Interpreter
else:
    from tensorflow.lite.python.interpreter import Interpreter


class TensorflowLiteDetection2D:

    # 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()

                dirModel = authenticationData['Configuration']['dir-model']
                graphModel = authenticationData['Configuration']['graph-model']
                labelMap = authenticationData['Configuration']['label-map']
                threshold = authenticationData['Configuration']['threshold']
                imageWidth = authenticationData['Configuration']['image-width']
                imageHeight = authenticationData['Configuration']['image-height']

                print("Dir model: " + str(dirModel))
                print("Graph model: " + str(graphModel))
                print("Label map: " + str(labelMap))
                print("Threshold: " + str(threshold))
                print("Image width: " + str(imageWidth))
                print("Image height: " + str(imageHeight))

                # Adjust threshold like float
                threshold = float(threshold)

                # 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 dirModel, graphModel, labelMap, threshold, imageWidth, imageHeight

    # Function: getLabels
    def getLabels(self, dirModel, labelMap):

        # Prepare label map path
        labelMapPath = str(dirModel) + "/" + str(labelMap)

        # Read label maps
        with open(labelMapPath, 'r') as file:
            labels = [line.strip() for line in file.readlines()]

        if labels[0] == '???':
            del (labels[0])

        return labels

    # Function: getModel
    def getModel(self, dirModel, graphModel):

        # Prepare modelPath
        modelPath = str(dirModel) + "/" + str(graphModel)

        # Load model
        model = Interpreter(model_path=modelPath)
        model.allocate_tensors()

        # Get model details
        inputDetails = model.get_input_details()
        outputDetails = model.get_output_details()
        floatingModel = (inputDetails[0]['dtype'] == np.float32)
        inputMean = 127.5
        inputSTD = 127.5

        return model, inputDetails, outputDetails, floatingModel, inputMean, inputSTD

    # Function: analyzeImage
    def analyzeImage(self, dataToSolve, model, inputDetails, outputDetails, floatingModel, inputMean, inputSTD, labels, threshold, imageWidth, imageHeight, outputDataPort):

        # Get tensor size parameters
        tensorHeight = inputDetails[0]['shape'][1]
        tensorWidth = inputDetails[0]['shape'][2]

        # Resize image to be processed in tensor
        dataToSolveTensor = cv2.resize(dataToSolve, (tensorWidth, tensorHeight))

        # Numpy expand sahpe
        inputData = np.expand_dims(dataToSolveTensor, axis=0)

        # Normalize pixel values if using a floating model
        if floatingModel:
            inputData = (np.float32(inputData) - inputMean) / inputSTD

        # Set tensor
        model.set_tensor(inputDetails[0]['index'], inputData)
        model.invoke()

        # Get results from tensor
        boxes = model.get_tensor(outputDetails[0]['index'])[0]
        classes = model.get_tensor(outputDetails[1]['index'])[0]
        scores = model.get_tensor(outputDetails[2]['index'])[0]

        for i in range(len(scores)):

            # If high level of certain
            if (scores[i] > threshold) and (scores[i] <= 1.0):

                # Get rectangle square coordinates
                yMin = int(max(1, (boxes[i][0] * int(imageHeight))))
                xMin = int(max(1, (boxes[i][1] * int(imageWidth))))
                yMax = int(min(int(imageHeight), (boxes[i][2] * int(imageHeight))))
                xMax = int(min(int(imageWidth), (boxes[i][3] * int(imageWidth))))

                # Draw rectangle of detected object
                cv2.rectangle(dataToSolve, (xMin, yMin), (xMax, yMax), (10, 255, 0), 2)

                # Get detected object label and score
                detectionObjectName = labels[int(classes[i])]
                detectionScore = int(scores[i] * 100)
                detectedObjectLabel = str(detectionObjectName) + ": " + str(detectionScore) + "%"

                # Get label size to be draw
                labelSize, baseLine = cv2.getTextSize(detectedObjectLabel, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
                labelYMin = max(yMin, labelSize[1] + 10)

                # Draw rectangle that will contain detected object
                cv2.rectangle(dataToSolve, (xMin, labelYMin - labelSize[1] - 10), (xMin + labelSize[0], labelYMin + baseLine - 10), (255, 255, 255), cv2.FILLED)

                # Draw detected class label
                cv2.putText(dataToSolve, detectedObjectLabel, (xMin, labelYMin - 7), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)

                # Get coordinates
                coordinatesXY = self.getCoordinates(xMin, labelSize, imageHeight, labelYMin, baseLine)

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

                # Send detected results
                outputDataPort.send(dataSolvedResults)

            elif scores[0] < 0.5:

                dataSolvedResults = "Detection: None, Score: 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, labelSize, imageHeight, labelYMin, baseLine):

        # Get centroid coordinates
        x = (xMin + xMin + labelSize[0]) / 2
        y = int(imageHeight) - labelYMin + baseLine - 10

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

        return coordinatesXY

    # Function: processRequest
    def processRequests(self, model, inputDetails, outputDetails, floatingModel, inputMean, inputSTD, labels, threshold, 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, model, inputDetails, outputDetails, floatingModel, inputMean, inputSTD, labels, threshold, imageWidth, imageHeight, 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: Tensorflow Lite Detection 2D                     ")
    print("                     Author: David Velasco Garcia                         ")
    print("                             @davidvelascogarcia                          ")
    print("**************************************************************************")
    print("**************************************************************************")

    print("\nLoading Tensorflow Lite Detection 2D engine ...\n")

    # Build tensorflowLiteDetection2D object
    tensorflowLiteDetection2D = TensorflowLiteDetection2D()

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

    # Get authentication data
    dirModel, graphModel, labelMap, threshold, imageWidth, imageHeight = tensorflowLiteDetection2D.getAuthenticationData()

    # Get labels
    labels = tensorflowLiteDetection2D.getLabels(dirModel, labelMap)

    # Get model
    model, inputDetails, outputDetails, floatingModel, inputMean, inputSTD = tensorflowLiteDetection2D.getModel(dirModel, graphModel)

    # Init Yarp network
    yarp.Network.init()

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

    # Process input requests
    tensorflowLiteDetection2D.processRequests(model, inputDetails, outputDetails, floatingModel, inputMean, inputSTD, labels, threshold, imageWidth, imageHeight, inputImagePort, outputImagePort, outputDataPort)

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

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


if __name__ == "__main__":

    # Call main function
    main()