///*************************************************************************************
//
// File name: ClassSurfaceSteel.cpp
// Location: See Matrox Example Launcher in the MIL Control Center
// 
//
// Synopsis:  This example detects the presence of scratch and 
// pit defects on the surface steel. 
//
// This example and data was inspired from the work of:
// Kechen Song and Yunhui Yan, "Micro surface defect detection method 
// for silicon steel strip based on saliency convex active contour model".
//
// Copyright (C) Matrox Electronic Systems Ltd., 1992-2020.
// All Rights Reserved

#include <mil.h>

// Path definitions.
#define EXAMPLE_IMAGE_DIR_PATH      M_IMAGE_PATH MIL_TEXT("/Classification/SurfaceSteel/")
#define EXAMPLE_CLASS_CTX_PATH      EXAMPLE_IMAGE_DIR_PATH MIL_TEXT("MatroxNet_SurfaceSteel.mclass")
#define TARGET_IMAGE_DIR_PATH       EXAMPLE_IMAGE_DIR_PATH MIL_TEXT("Products")

// Util constant.
#define BUFFERING_SIZE_MAX 10

// Function declarations.
struct ClassStruct
{
   MIL_INT NbCategories,
           NbOfFrames;

   MIL_ID ClassCtx,
          ClassRes,
          MilDisplay,
          MilDispChild;
   };

MIL_INT MFTYPE ClassificationFunc(MIL_INT HookType,
                                  MIL_ID  EventId,
                                  void*   pHookData);

void SetupDisplay(MIL_ID  MilSystem,
                  MIL_ID  MilDisplay,
                  MIL_INT SourceSizeX,
                  MIL_INT SourceSizeY,
                  MIL_ID  ClassCtx,
                  MIL_ID  &MilDispImage,
                  MIL_ID  &MilDispChild,
                  MIL_ID  &MilOverlay,
                  MIL_INT NbCategories);

#define SYSTEM_TO_USE M_SYSTEM_HOST
#define DCF_TO_USE    TARGET_IMAGE_DIR_PATH

// ****************************************************************************
//    Main.
// ****************************************************************************
int main(void)
   {
   MIL_ID MilApplication,    // MIL application identifier
          MilSystem,         // MIL system identifier
          MilDisplay,        // MIL display identifier
          MilOverlay,        // MIL overlay identifier
          MilDigitizer,      // MIL digitizer identifier
          MilDispImage,      // MIL image identifier
          MilDispChild,      // MIL image identifier
          ClassCtx,          // MIL classification Context
          ClassRes;          // MIL classification Result

   MIL_ID MilGrabBufferList[BUFFERING_SIZE_MAX]; // MIL image identifier

   MIL_INT NumberOfCategories,
           BufIndex,
           SourceSizeX,
           SourceSizeY,
           InputSizeX,
           InputSizeY;

   // Allocate MIL objects.
   MappAlloc(M_NULL, M_DEFAULT, &MilApplication);
   MsysAlloc(M_DEFAULT, SYSTEM_TO_USE, M_DEFAULT, M_DEFAULT, &MilSystem);
   MdispAlloc(MilSystem, M_DEFAULT, MIL_TEXT("M_DEFAULT"), M_DEFAULT, &MilDisplay);
   MdigAlloc(MilSystem, M_DEFAULT, DCF_TO_USE, M_DEFAULT, &MilDigitizer);
   
   // Print the example synopsis.
   MosPrintf(MIL_TEXT("[EXAMPLE NAME]\n"));
   MosPrintf(MIL_TEXT("ClassSurfaceSteel\n\n"));
   MosPrintf(MIL_TEXT("[SYNOPSIS]\n"));
   MosPrintf(MIL_TEXT("This programs shows the use of a pre-trained classification context\n"));
   MosPrintf(MIL_TEXT("to detect the presence of defects on surface steel.\n\n"));
   MosPrintf(MIL_TEXT("[MODULES USED]\n"));
   MosPrintf(MIL_TEXT("Classification, Buffer, Display, Graphics, Image Processing.\n\n"));

   // Wait for user.
   MosPrintf(MIL_TEXT("Press <Enter> to continue.\n"));
   MosGetchar();
   
   // Restore the trained context.
   MosPrintf(MIL_TEXT("Restoring the classification context from file.."));
   MclassRestore(EXAMPLE_CLASS_CTX_PATH, MilSystem, M_DEFAULT, &ClassCtx);
   MosPrintf(MIL_TEXT("."));

   // Preprocess the context.
   MclassPreprocess(ClassCtx, M_DEFAULT);
   MosPrintf(MIL_TEXT(".ready.\n"));

   MclassInquire(ClassCtx, M_CONTEXT, M_NUMBER_OF_CLASSES   + M_TYPE_MIL_INT, &NumberOfCategories);
   MclassInquire(ClassCtx, M_DEFAULT_SOURCE_LAYER, M_SIZE_X + M_TYPE_MIL_INT, &InputSizeX);
   MclassInquire(ClassCtx, M_DEFAULT_SOURCE_LAYER, M_SIZE_Y + M_TYPE_MIL_INT, &InputSizeY);

   // Inquire and print source layer information.
   MosPrintf(MIL_TEXT(" - The classifier was trained to recognize %d categories\n"), NumberOfCategories);
   MosPrintf(MIL_TEXT(" - The classifier was trained for %dx%d source images\n\n"), InputSizeX, InputSizeY);

   // Inquire the size of the source image.
   MdigInquire(MilDigitizer, M_SIZE_X, &SourceSizeX);
   MdigInquire(MilDigitizer, M_SIZE_Y, &SourceSizeY);

   // Set the expected target image sizes.
   MclassControl(ClassCtx, M_CONTEXT, M_TARGET_IMAGE_SIZE_X, SourceSizeX);
   MclassControl(ClassCtx, M_CONTEXT, M_TARGET_IMAGE_SIZE_Y, SourceSizeY);
   
   // Preprocess the context.
   MclassPreprocess(ClassCtx, M_DEFAULT);

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

   // Retrieve the number of frame in the source directory.
   MIL_INT NumberOfFrames = 0;
   MdigInquire(MilDigitizer, M_SOURCE_NUMBER_OF_FRAMES, &NumberOfFrames);

   // Allocate a classification result buffer.
   MclassAllocResult(MilSystem, M_PREDICT_CNN_RESULT, M_DEFAULT, &ClassRes);

   // Prepare data for Hook Function.
   ClassStruct ClassificationData;
   ClassificationData.ClassCtx        = ClassCtx;
   ClassificationData.ClassRes        = ClassRes;
   ClassificationData.MilDisplay      = MilDisplay;
   ClassificationData.MilDispChild    = MilDispChild; 
   ClassificationData.NbCategories    = NumberOfCategories;
   ClassificationData.NbOfFrames      = NumberOfFrames;

   // Allocate the grab buffers.
   for(BufIndex = 0; BufIndex<BUFFERING_SIZE_MAX; BufIndex++)
      { MbufAllocColor(MilSystem, 3, SourceSizeX, SourceSizeY, 8 + M_UNSIGNED, M_IMAGE + M_GRAB + M_PROC, &MilGrabBufferList[BufIndex]); }

   // Start the grab.
   if (NumberOfFrames != M_INFINITE)
      { MdigProcess(MilDigitizer, MilGrabBufferList, BUFFERING_SIZE_MAX, M_SEQUENCE + M_COUNT(NumberOfFrames), M_SYNCHRONOUS, &ClassificationFunc, &ClassificationData); }
   else
      { MdigProcess(MilDigitizer, MilGrabBufferList, BUFFERING_SIZE_MAX, M_START, M_DEFAULT, &ClassificationFunc, &ClassificationData); }
   
   // Ready to exit.
   MosPrintf(MIL_TEXT("\nPress <Enter> to exit.\n"));
   MosGetch();

   // Stop the digitizer.
   MdigProcess(MilDigitizer, MilGrabBufferList, BUFFERING_SIZE_MAX, M_STOP, M_DEFAULT, M_NULL, M_NULL);

   // Free the allocated resources.
   MdigFree(MilDigitizer);
   MbufFree(MilDispChild);
   MbufFree(MilDispImage);

   for(BufIndex = 0; BufIndex < BUFFERING_SIZE_MAX; BufIndex++)
      {
      MbufFree(MilGrabBufferList[BufIndex]);
      MilGrabBufferList[BufIndex] = M_NULL;
      }

   MclassFree(ClassRes);
   MclassFree(ClassCtx);
   MdispFree(MilDisplay);
   MsysFree(MilSystem);
   MappFree(MilApplication);

   return 0;
   }

void SetupDisplay(MIL_ID  MilSystem,
                  MIL_ID  MilDisplay,
                  MIL_INT SourceSizeX,
                  MIL_INT SourceSizeY,
                  MIL_ID  ClassCtx,
                  MIL_ID  &MilDispImage,
                  MIL_ID  &MilDispChild,
                  MIL_ID  &MilOverlay,
                  MIL_INT NbCategories
                  )
   {

   MIL_ID MilImageLoader,  // MIL image identifier       
          MilChildSample;  // MIL child image identifier
          
   // Allocate a color buffer.
   MIL_INT IconSize = SourceSizeY / NbCategories;
   MilDispImage = MbufAllocColor(MilSystem, 3, SourceSizeX + IconSize, SourceSizeY, 8 + M_UNSIGNED, M_IMAGE + M_PROC + M_DISP, M_NULL);
   MilDispChild = MbufChild2d(MilDispImage, 0, 0, SourceSizeX, SourceSizeY, M_NULL);
   MbufClear(MilDispImage, M_COLOR_BLACK);

   // Set annotation color.
   MgraColor(M_DEFAULT, M_COLOR_RED);

   // Setup the display.
   for(int iter = 0; iter < NbCategories; iter++)
      {
      // Allocate a child buffer per product categorie.   
      MbufChild2d(MilDispImage, SourceSizeX, iter * IconSize, IconSize, IconSize, &MilChildSample);
      MbufClear(MilChildSample, M_COLOR_BLACK);

      // Load the sample image.
      MclassInquire(ClassCtx, M_CLASS_INDEX(iter), M_CLASS_ICON_ID + M_TYPE_MIL_ID, &MilImageLoader);

      // Retrieve the class description.
      MIL_STRING Text;
      MclassInquire(ClassCtx, M_CLASS_INDEX(iter), M_CLASS_NAME, Text);
      
      if (MilImageLoader != M_NULL)
         {
         // Retrieve the color associated to the class.
         MIL_DOUBLE ClassColor;
         MclassInquire(ClassCtx, M_CLASS_INDEX(iter), M_CLASS_DRAW_COLOR, &ClassColor);
         
         // Draw the class name using the color associated to the class.
         MgraColor(M_DEFAULT, ClassColor);
         MgraText(M_DEFAULT, MilChildSample, 10, 10, Text);

         MIL_INT ClassImageExampleSizeX = MbufInquire(MilImageLoader, M_SIZE_X, M_NULL);
         MIL_INT ClassImageExampleSizeY = MbufInquire(MilImageLoader, M_SIZE_Y, M_NULL);

         if ((ClassImageExampleSizeX >= IconSize) || (ClassImageExampleSizeY >= IconSize))
            {
            MimResize(MilImageLoader, MilChildSample, M_FILL_DESTINATION, M_FILL_DESTINATION, M_AVERAGE);
            }
         else
            {
            MIL_INT OffsetX = (IconSize - ClassImageExampleSizeX) / 2;
            MIL_INT OffsetY = (IconSize - ClassImageExampleSizeY) / 2;
            MbufCopyColor2d(MilImageLoader, MilChildSample, M_ALL_BANDS, 0, 0, M_ALL_BANDS, OffsetX, OffsetY, ClassImageExampleSizeX, ClassImageExampleSizeY);
            }
         }

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

      // Free the allocated buffers.
      MbufFree(MilChildSample);
      }

   // Display the window with black color.
   MdispSelect(MilDisplay, MilDispImage);

   // Prepare for overlay annotations.
   MdispControl(MilDisplay, M_OVERLAY, M_ENABLE);
   MilOverlay = MdispInquire(MilDisplay, M_OVERLAY_ID, M_NULL);
   }

MIL_INT MFTYPE ClassificationFunc(MIL_INT /*HookType*/, MIL_ID EventId, void* DataPtr)
   {
   MIL_ID MilImage;
   MdigGetHookInfo(EventId, M_MODIFIED_BUFFER + M_BUFFER_ID, &MilImage);

   ClassStruct* data = static_cast<ClassStruct*>(DataPtr);
  
   // Perform product recognition using the classification module.
   MclassPredict(data->ClassCtx, MilImage, data->ClassRes, M_DEFAULT);

   // Clear the overlay buffer.
   MdispControl(data->MilDisplay, M_OVERLAY_CLEAR, M_TRANSPARENT_COLOR);

   // Inquire the MIL_ID of the display overlay.
   MIL_ID DisplayOverlayId = M_NULL;
   MdispInquire(data->MilDisplay, M_OVERLAY_ID, &DisplayOverlayId);

   // Display the new target image.
   MbufCopy(MilImage, data->MilDispChild); 
   for(int i = 1; i < data->NbCategories; i++)
      { MclassDraw(M_DEFAULT, data->ClassRes, DisplayOverlayId, M_DRAW_BEST_INDEX_CONTOUR_IMAGE + M_PSEUDO_COLOR, M_CLASS_INDEX(i), M_DEFAULT); }
   
   // Wait for the user.
   if (data->NbOfFrames != M_INFINITE)
      {
      MosPrintf(MIL_TEXT("Press <Enter> to continue.\r"));
      MosGetch();
      }

   return 0;
   }