' File name: Mclass.vb
' Location: See Matrox Example Launcher in the MIL Control Center
' Synopsis:  This example identifies the type of pastas using a 
' pre-trained classification module. 
' Copyright (C) Matrox Electronic Systems Ltd., 1992-2020.
' All Rights Reserved
Imports Microsoft.VisualBasic
Imports System

Imports Matrox.MatroxImagingLibrary
Imports System.Runtime.InteropServices

Namespace Mclass
    ' Processing function parameters structure.
    Public Class ClassStruct
        Public NbCategories As MIL_INT
        Public NbOfFrames As MIL_INT
        Public SourceSizeX As MIL_INT
        Public SourceSizeY As MIL_INT

        Public ClassCtx As MIL_ID
        Public ClassRes As MIL_ID
        Public MilDisplay As MIL_ID
        Public MilDispChild As MIL_ID
        Public MilOverlayImage As MIL_ID
    End Class

    Public Class Program
        ' Path definitions.
        Private Const EXAMPLE_IMAGE_DIR_PATH As String = MIL.M_IMAGE_PATH & "/Classification/Pasta/"
        Private Const EXAMPLE_CLASS_CTX_PATH As String = EXAMPLE_IMAGE_DIR_PATH & "MatroxNet_PastaEx.Mclass"
        Private Const TARGET_IMAGE_DIR_PATH As String = EXAMPLE_IMAGE_DIR_PATH & "Products"

        ' Util constant.
        Private Const BUFFERING_SIZE_MAX As Integer = 10

        ' Use the images from the example folder by default.
        ' Add the USE_EXAMPLE_IMAGE_FOLDER to the predefined constants in project properties to change this.
        Private Const SYSTEM_TO_USE As String = MIL.M_SYSTEM_HOST
        Private Const DCF_TO_USE As String = TARGET_IMAGE_DIR_PATH
        Private Const SYSTEM_TO_USE As String = MIL.M_SYSTEM_DEFAULT
        Private Const DCF_TO_USE As String = "M_DEFAULT"
#End If

        Shared Function Main(ByVal args() As String) As Integer
            Dim MilApplication As MIL_ID = MIL.M_NULL ' MIL application identifier
            Dim MilSystem As MIL_ID = MIL.M_NULL ' MIL system identifier
            Dim MilDisplay As MIL_ID = MIL.M_NULL ' MIL display identifier
            Dim MilOverlay As MIL_ID = MIL.M_NULL ' MIL overlay identifier
            Dim MilDigitizer As MIL_ID = MIL.M_NULL ' MIL digitizer identifier
            Dim MilDispImage As MIL_ID = MIL.M_NULL ' MIL image identifier
            Dim MilDispChild As MIL_ID = MIL.M_NULL ' MIL image identifier
            Dim ClassCtx As MIL_ID = MIL.M_NULL ' MIL classification Context
            Dim ClassRes As MIL_ID = MIL.M_NULL ' MIL classification Result

            Dim MilGrabBufferList(BUFFERING_SIZE_MAX - 1) As MIL_ID ' MIL image identifier
            Dim MilChildBufferList(BUFFERING_SIZE_MAX - 1) As MIL_ID ' MIL child identifier

            Dim NumberOfCategories As MIL_INT = 0
            Dim BufIndex As Integer = 0
            Dim SourceSizeX As MIL_INT = 0
            Dim SourceSizeY As MIL_INT = 0
            Dim InputSizeX As MIL_INT = 0
            Dim InputSizeY As MIL_INT = 0

            ' Allocate MIL objects.
            MIL.MappAlloc(MIL.M_NULL, MIL.M_DEFAULT, MilApplication)
            MIL.MdispAlloc(MilSystem, MIL.M_DEFAULT, "MIL.M_DEFAULT", MIL.M_DEFAULT, MilDisplay)
            MIL.MdigAlloc(MilSystem, MIL.M_DEFAULT, DCF_TO_USE, MIL.M_DEFAULT, MilDigitizer)

            ' Print the example synopsis.
            Console.WriteLine("[EXAMPLE NAME]")
            Console.WriteLine("This programs shows the use of a pre-trained classification")
            Console.WriteLine("tool to recognize product categories.")
            Console.WriteLine("[MODULES USED]")
            Console.WriteLine("Classification, Buffer, Display, Graphics, Image Processing.")

            ' Wait for user.
            Console.WriteLine("Press <Enter> to continue.")

            Console.Write("Restoring the classification context from file..")
            MIL.MclassRestore(EXAMPLE_CLASS_CTX_PATH, MilSystem, MIL.M_DEFAULT, ClassCtx)

            ' Preprocess the context.
            MIL.MclassPreprocess(ClassCtx, MIL.M_DEFAULT)

            MIL.MclassInquire(ClassCtx, MIL.M_CONTEXT, MIL.M_NUMBER_OF_CLASSES + MIL.M_TYPE_MIL_INT, NumberOfCategories)
            MIL.MclassInquire(ClassCtx, MIL.M_DEFAULT_SOURCE_LAYER, MIL.M_SIZE_X + MIL.M_TYPE_MIL_INT, InputSizeX)
            MIL.MclassInquire(ClassCtx, MIL.M_DEFAULT_SOURCE_LAYER, MIL.M_SIZE_Y + MIL.M_TYPE_MIL_INT, InputSizeY)

            If (NumberOfCategories > 0) Then
                ' Inquire and print source layer information.
                Console.WriteLine(" - The classifier was trained to recognize {0} categories", NumberOfCategories)
                Console.WriteLine(" - The classifier was trained for {0}x{1} source images", InputSizeX, InputSizeY)

                ' Allocate a classification result buffer.
                MIL.MclassAllocResult(MilSystem, MIL.M_PREDICT_CNN_RESULT, MIL.M_DEFAULT, ClassRes)

                ' Inquire the size of the source image.
                MIL.MdigInquire(MilDigitizer, MIL.M_SIZE_X, SourceSizeX)
                MIL.MdigInquire(MilDigitizer, MIL.M_SIZE_Y, SourceSizeY)

                ' Setup the example display.
                SetupDisplay(MilSystem, MilDisplay, SourceSizeX, SourceSizeY, ClassCtx, MilDispImage, MilDispChild, MilOverlay, NumberOfCategories)

                ' Retrieve the number of frame in the source directory.
                Dim NumberOfFrames As MIL_INT = 0
                MIL.MdigInquire(MilDigitizer, MIL.M_SOURCE_NUMBER_OF_FRAMES, NumberOfFrames)

                ' Prepare data for Hook Function.
                Dim ClassificationData As New ClassStruct()
                ClassificationData.ClassCtx = ClassCtx
                ClassificationData.ClassRes = ClassRes
                ClassificationData.MilDisplay = MilDisplay
                ClassificationData.MilDispChild = MilDispChild
                ClassificationData.NbCategories = NumberOfCategories
                ClassificationData.MilOverlayImage = MilOverlay
                ClassificationData.SourceSizeX = SourceSizeX
                ClassificationData.SourceSizeY = SourceSizeY
                ClassificationData.NbOfFrames = NumberOfFrames

                ' Allocate the grab buffers.
                For BufIndex = 0 To BUFFERING_SIZE_MAX - 1
                    MIL.MbufAlloc2d(MilSystem, SourceSizeX, SourceSizeY, 8 + MIL.M_UNSIGNED, MIL.M_IMAGE + MIL.M_GRAB + MIL.M_PROC, MilGrabBufferList(BufIndex))
                    MIL.MbufChild2d(MilGrabBufferList(BufIndex), (SourceSizeX - InputSizeX) / 2, (SourceSizeY - InputSizeY) / 2, InputSizeX, InputSizeY, MilChildBufferList(BufIndex))
                    MIL.MobjControl(MilGrabBufferList(BufIndex), MIL.M_OBJECT_USER_DATA_PTR, MilChildBufferList(BufIndex))
                Next BufIndex

                ' Start the grab.
                Dim ClassificationFuncHook As New MIL_DIG_HOOK_FUNCTION_PTR(AddressOf ClassificationFunc)
                Dim ClassificationDataHandle As GCHandle = GCHandle.Alloc(ClassificationData)
                If NumberOfFrames <> MIL.M_INFINITE Then
                    MIL.MdigProcess(MilDigitizer, MilGrabBufferList, BUFFERING_SIZE_MAX, MIL.M_SEQUENCE + MIL.M_COUNT(NumberOfFrames), MIL.M_SYNCHRONOUS, ClassificationFuncHook, GCHandle.ToIntPtr(ClassificationDataHandle))
                    MIL.MdigProcess(MilDigitizer, MilGrabBufferList, BUFFERING_SIZE_MAX, MIL.M_START, MIL.M_DEFAULT, ClassificationFuncHook, GCHandle.ToIntPtr(ClassificationDataHandle))
                End If

                ' Ready to exit.
                Console.WriteLine("Press <Enter> to exit.")

                ' Stop the digitizer.
                MIL.MdigProcess(MilDigitizer, MilGrabBufferList, BUFFERING_SIZE_MAX, MIL.M_STOP, MIL.M_DEFAULT, MIL.M_NULL, MIL.M_NULL)


                ' Free the allocated resources.                

                For BufIndex = 0 To BUFFERING_SIZE_MAX - 1
                Next BufIndex

            End If



            Return 0
        End Function

        Private Shared Sub SetupDisplay(ByVal MilSystem As MIL_ID,
                                        ByVal MilDisplay As MIL_ID,
                                        ByVal SourceSizeX As MIL_INT,
                                        ByVal SourceSizeY As MIL_INT,
                                        ByVal ClassCtx As MIL_ID,
                                        ByRef MilDispImage As MIL_ID,
                                        ByRef MilDispChild As MIL_ID,
                                        ByRef MilOverlay As MIL_ID,
                                        ByVal NbCategories As MIL_INT)
            Dim MilImageLoader As MIL_ID = MIL.M_NULL ' MIL image identifier
            Dim MilChildSample As MIL_ID = MIL.M_NULL ' MIL child image identifier

            ' Allocate a color buffer.
            Dim IconSize As MIL_INT = SourceSizeY / NbCategories
            MilDispImage = MIL.MbufAllocColor(MilSystem, 3, SourceSizeX + IconSize, SourceSizeY, 8 + MIL.M_UNSIGNED, MIL.M_IMAGE + MIL.M_PROC + MIL.M_DISP, MIL.M_NULL)
            MIL.MbufClear(MilDispImage, MIL.M_COLOR_BLACK)
            MilDispChild = MIL.MbufChild2d(MilDispImage, 0, 0, SourceSizeX, SourceSizeY, MIL.M_NULL)

            ' Set annotation color.
            MIL.MgraColor(MIL.M_DEFAULT, MIL.M_COLOR_RED)

            ' Setup the display.
            For iter As Integer = 0 To CInt(NbCategories) - 1
                ' Allocate a child buffer per product categorie.   
                MIL.MbufChild2d(MilDispImage, SourceSizeX, iter * IconSize, IconSize, IconSize, MilChildSample)

                ' Load the sample image.
                MIL.MclassInquire(ClassCtx, MIL.M_CLASS_INDEX(iter), MIL.M_CLASS_ICON_ID + MIL.M_TYPE_MIL_ID, MilImageLoader)

                If MilImageLoader <> MIL.M_NULL Then
                End If

                ' Draw an initial red rectangle around the buffer.
                MIL.MgraRect(MIL.M_DEFAULT, MilChildSample, 0, 1, IconSize - 1, IconSize - 2)

                ' Free the allocated buffers.
            Next iter

            ' Display the window with black color.
            MIL.MdispSelect(MilDisplay, MilDispImage)

            ' Prepare for overlay annotations.
            MIL.MdispControl(MilDisplay, MIL.M_OVERLAY, MIL.M_ENABLE)
            MilOverlay = CType(MIL.MdispInquire(MilDisplay, MIL.M_OVERLAY_ID, MIL.M_NULL), MIL_ID)
        End Sub

        Private Shared Function ClassificationFunc(ByVal HookType As MIL_INT, ByVal EventId As MIL_ID, ByVal DataPtr As IntPtr) As MIL_INT
            Dim MilImage As MIL_ID = MIL.M_NULL
            Dim pMilInputImage As MIL_ID = MIL.M_NULL

            MIL.MdigGetHookInfo(EventId, MIL.M_MODIFIED_BUFFER + MIL.M_BUFFER_ID, MilImage)

            Dim data As ClassStruct = CType(GCHandle.FromIntPtr(DataPtr).Target, ClassStruct)
            MIL.MdispControl(data.MilDisplay, MIL.M_UPDATE, MIL.M_DISABLE)
            pMilInputImage = CType(MIL.MobjInquire(MilImage, MIL.M_OBJECT_USER_DATA_PTR, CType(MIL.M_NULL, IntPtr)), MIL_ID)

            ' Display the new target image.
            MIL.MbufCopy(MilImage, data.MilDispChild)

            ' Perform product recognition using the classification module.
            MIL.MclassPredict(data.ClassCtx, pMilInputImage, data.ClassRes, MIL.M_DEFAULT)

            ' Retrieve best classification score and class index.
            Dim BestScore As Double = 0
            MIL.MclassGetResult(data.ClassRes, MIL.M_GENERAL, MIL.M_BEST_CLASS_SCORE + MIL.M_TYPE_MIL_DOUBLE, BestScore)

            Dim BestIndex As MIL_INT = 0
            MIL.MclassGetResult(data.ClassRes, MIL.M_GENERAL, MIL.M_BEST_CLASS_INDEX + MIL.M_TYPE_MIL_INT, BestIndex)

            ' Clear the overlay buffer.
            MIL.MdispControl(data.MilDisplay, MIL.M_OVERLAY_CLEAR, MIL.M_TRANSPARENT_COLOR)

            ' Draw a green rectangle around the winning sample.
            Dim IconSize As MIL_INT = data.SourceSizeY / data.NbCategories
            MIL.MgraRect(MIL.M_DEFAULT, data.MilOverlayImage, data.SourceSizeX, (BestIndex * IconSize) + 1, data.SourceSizeX + IconSize - 1, (BestIndex + 1) * IconSize - 2)

            ' Print the classification accuracy in the sample buffer.
            Dim Accuracy_text As String = String.Format("{0:N1}% score", BestScore)
            MIL.MgraText(MIL.M_DEFAULT, data.MilOverlayImage, data.SourceSizeX + 2, BestIndex * IconSize + 4, Accuracy_text)

            ' Update the display.
            MIL.MdispControl(data.MilDisplay, MIL.M_UPDATE, MIL.M_ENABLE)

            ' Wait for the user.
            If data.NbOfFrames <> MIL.M_INFINITE Then
                Console.Write("Press <Enter> to continue." & Constants.vbCr)
            End If

            Return 0
        End Function
    End Class
End Namespace