SighthoundRecognizerComponent.cpp

// Fill out your copyright notice in the Description page of Project Settings.


#include "SighthoundRecognizerComponent.h"
#include "AppleARKitBlueprintFunctionLibrary.h"
#include "AppleARKitSession.h"
#if PLATFORM_IOS
#import <Foundation/Foundation.h>
#import <CoreMedia/CoreMedia.h>
#import <CoreVideo/CoreVideo.h>
#import <CoreGraphics/CoreGraphics.h>
#import <UIKit/UIKit.h>
#import <Accelerate/Accelerate.h>
#import <Photos/Photos.h>
#import <VideoToolbox/VideoToolbox.h>
#include "SIOFramework/SIOFramework.h"
#include "RenderingThread.h"
#endif

#if PLATFORM_IOS

namespace
{
    vImage_Error configureYpCbCrToARGBInfo(vImage_YpCbCrToARGB& matrix)
    {
        auto pixelRange = vImage_YpCbCrPixelRange { .Yp_bias= 0,
            .CbCr_bias= 128,
            .YpRangeMax= 255,
            .CbCrRangeMax= 255,
            .YpMax= 255,
            .YpMin= 1,
            .CbCrMax= 255,
            .CbCrMin= 0};
        auto error = vImageConvert_YpCbCrToARGB_GenerateConversion(kvImage_YpCbCrToARGBMatrix_ITU_R_709_2, &pixelRange, &matrix, kvImage420Yp8_CbCr8, kvImageARGB8888, UInt32(kvImageNoFlags));
        
        return error;
    }
    
    
    
    CVPixelBufferRef Yuv420ToRgbImpl(CVPixelBufferRef input, vImage_Buffer& scaleBuffer, vImage_Buffer& rotationBuffer, vImage_Buffer& destinationBuffer)
    {
        auto cgImageFormat = vImage_CGImageFormat{.bitsPerComponent= 8,
            .bitsPerPixel= 32,
            .colorSpace=  nil,
            .bitmapInfo= kCGImageAlphaNoneSkipFirst,
            .version= 0,
            .decode= nil,
            .renderingIntent= kCGRenderingIntentDefault};
        auto pixelBuffer = input;
        CVPixelBufferLockBaseAddress(pixelBuffer, kCVPixelBufferLock_ReadOnly);
        auto lumaBaseAddress = CVPixelBufferGetBaseAddressOfPlane(pixelBuffer, 0);
        auto lumaWidth = CVPixelBufferGetWidthOfPlane(pixelBuffer, 0);
        auto lumaHeight = CVPixelBufferGetHeightOfPlane(pixelBuffer, 0);
        auto lumaRowBytes = CVPixelBufferGetBytesPerRowOfPlane(pixelBuffer, 0);
        
        auto sourceLumaBuffer = vImage_Buffer {.data= lumaBaseAddress,
            .height= vImagePixelCount(lumaHeight),
            .width= vImagePixelCount(lumaWidth),
            .rowBytes= lumaRowBytes} ;
        
        auto chromaBaseAddress = CVPixelBufferGetBaseAddressOfPlane(pixelBuffer, 1);
        auto chromaWidth = CVPixelBufferGetWidthOfPlane(pixelBuffer, 1);
        auto chromaHeight = CVPixelBufferGetHeightOfPlane(pixelBuffer, 1);
        auto chromaRowBytes = CVPixelBufferGetBytesPerRowOfPlane(pixelBuffer, 1);
        
        auto sourceChromaBuffer = vImage_Buffer {.data= chromaBaseAddress,
            .height= vImagePixelCount(chromaHeight),
            .width= vImagePixelCount(chromaWidth),
            .rowBytes= chromaRowBytes} ;
        if (scaleBuffer.data == nil)
        {
            auto error = vImageBuffer_Init(&scaleBuffer,
                                           sourceLumaBuffer.height,
                                           sourceLumaBuffer.width,
                                           cgImageFormat.bitsPerPixel,
                                           vImage_Flags(kvImageNoFlags));
        }
        if (rotationBuffer.data == nil)
        {
            auto error = vImageBuffer_Init(&rotationBuffer,
                                           scaleBuffer.height/2,
                                           scaleBuffer.width/2,
                                           cgImageFormat.bitsPerPixel,
                                           vImage_Flags(kvImageNoFlags));
        }
        if (destinationBuffer.data == nil)
        {
            auto error = vImageBuffer_Init(&destinationBuffer,
                                           rotationBuffer.width,
                                           rotationBuffer.height,
                                           cgImageFormat.bitsPerPixel,
                                           vImage_Flags(kvImageNoFlags));
        }
        auto infoYpCbCrToARGB = vImage_YpCbCrToARGB();
        auto error = configureYpCbCrToARGBInfo(infoYpCbCrToARGB);
        uint8_t permuteMap[4] = {3, 2, 1, 0};
        error = vImageConvert_420Yp8_CbCr8ToARGB8888(&sourceLumaBuffer,
                                                     &sourceChromaBuffer,
                                                     &scaleBuffer,
                                                     &infoYpCbCrToARGB,
                                                     permuteMap,
                                                     255,
                                                     vImage_Flags(kvImagePrintDiagnosticsToConsole));
        CVPixelBufferUnlockBaseAddress(pixelBuffer, kCVPixelBufferLock_ReadOnly);
        error = vImageScale_ARGB8888(&scaleBuffer, &rotationBuffer, NULL, kvImagePrintDiagnosticsToConsole);
        error = vImageRotate90_ARGB8888(&rotationBuffer, &destinationBuffer, kRotate90DegreesClockwise, Pixel_8888{0,0,0,0},vImage_Flags(kvImagePrintDiagnosticsToConsole));
        
        static bool bSaved = false;
        if (!bSaved)
        {
            bSaved = true;
            vImage_CGImageFormat format = {
                .bitsPerComponent = 8,
                .bitsPerPixel = 32,
                .colorSpace = NULL,
                .bitmapInfo = (CGBitmapInfo)(kCGImageAlphaFirst|kCGBitmapByteOrder32Little),
                .version = 0,
                .decode = NULL,
                .renderingIntent = kCGRenderingIntentDefault,
            };
            CGImageRef destRef = vImageCreateCGImageFromBuffer( &destinationBuffer, &format, NULL, NULL, kvImageNoFlags, &error);
            UIImage* Image = [[UIImage alloc] initWithCGImage:destRef];
            CGImageRelease(destRef);
            [[PHPhotoLibrary sharedPhotoLibrary] performChanges:^{
                [PHAssetChangeRequest creationRequestForAssetFromImage:Image];
            } completionHandler:^(BOOL success, NSError *error) {
                if (success) {
                    NSLog(@"Successfully copied image to photos folder");
                }
                else {
                    NSLog(@"Failed to copy imge to photos folder : %@",error);
                }
                CFRelease(Image);
            }];
        }
        NSDictionary *options =
        [NSDictionary dictionaryWithObjectsAndKeys:
          [NSNumber numberWithBool:YES],
          kCVPixelBufferCGImageCompatibilityKey,
          [NSNumber numberWithBool:YES],
          kCVPixelBufferCGBitmapContextCompatibilityKey,
          nil];
         CVPixelBufferRef outPixelBuffer;
        //CVPixelBufferCreateWithBytes(CFAllocatorRef allocator, size_t width, size_t height, OSType pixelFormatType, void *baseAddress, size_t bytesPerRow, CVPixelBufferReleaseBytesCallback releaseCallback, void *releaseRefCon, CFDictionaryRef pixelBufferAttributes, CVPixelBufferRef  _Nullable *pixelBufferOut);
         CVReturn status =
            CVPixelBufferCreateWithBytes(
                                kCFAllocatorDefault, destinationBuffer.width, destinationBuffer.height,
                                kCVPixelFormatType_32BGRA, destinationBuffer.data, destinationBuffer.rowBytes, nullptr, nullptr, (__bridge CFDictionaryRef)options,
                                &outPixelBuffer);
        if (status != kCVReturnSuccess)
         {
             return nullptr;
         }
        return outPixelBuffer;
    }
    
    CVPixelBufferRef Yuv420ToRgb(CVPixelBufferRef input)
    {
        static bool bInit = false;
        static vImage_Buffer scaleBuffer;
        static vImage_Buffer rotationBuffer;
        static vImage_Buffer destinationBuffer;
        if (!bInit)
        {
            bInit = true;
            scaleBuffer.data = nullptr;
            rotationBuffer.data = nullptr;
            destinationBuffer.data = nullptr;
        }
        return Yuv420ToRgbImpl(input, scaleBuffer, rotationBuffer, destinationBuffer);
    }
}

void USighthoundRecognizerComponent::SetupPipeline()
{
    // Find the path to cars.yaml and load it into the config
    NSString *filePath = [[NSBundle mainBundle] pathForResource:@"cars" ofType:@"yaml"];
    UE_LOG(LogTemp, Log, TEXT("cars.yaml => %s"), *(FString(filePath)));
    SIOPipelineConfig *config = [[SIOPipelineConfig alloc] initWithConfig:filePath];
    // Get the path to where we will store the database file
    NSArray *paths = NSSearchPathForDirectoriesInDomains(NSDocumentDirectory, NSUserDomainMask, YES);
    NSString *documentsDirectory = [paths objectAtIndex:0];
   // NSMutableDictionary<NSString *, NSString *>* dict = [NSMutableDictionary dictionary];
    //[dict setValue:documentsDirectory forKey:@"DB_ROOT"];
    //[config applyParameters:dict];
    // Let it know what streams we are going to use
    [config exposeInputStream:@"input" buffer:-1];
    [config exposeOutputStream:@"boxes"];
    //[config exposeOutputStream:@"scores"];
   // [config exposeOutputStream:@"types"];
    [config exposeOutputStream:@"makemodels"];
    [config exposeOutputStream:@"colors"];
    Pipeline = [[SIOPipeline alloc] initWithConfig:config];
    [Pipeline start];
    UE_LOG(LogTemp, Log, TEXT("Started Pipeline"));
}


void USighthoundRecognizerComponent::RunRecognizer(CVPixelBufferRef pixelBuffer, TSharedPtr<TArray<FCarModelRecognition>, ESPMode::ThreadSafe> ResultPtr)
{
    if (Pipeline == nullptr) SetupPipeline();
    //UE_LOG(LogTemp, Log, TEXT("Writing input to pipeline %p"), pixelBuffer);
    CVPixelBufferLockBaseAddress(pixelBuffer, 0);
    [Pipeline writeInput:@"input" withCVPixelBufferRef: pixelBuffer];
    //UE_LOG(LogTemp, Log, TEXT("Wrote input to pipeline"));
    NSArray* BoxList = [Pipeline readOutputListCGRect:@"boxes"];
    NSArray* MakeModelList = [Pipeline readOutputListString:@"makemodels"];
    NSArray* ColorList = [Pipeline readOutputListString:@"colors"];
    CVPixelBufferUnlockBaseAddress(pixelBuffer, 0);
    auto Count = [BoxList count];
    //UE_LOG(LogTemp, Log, TEXT("Boxes: %d"), Count);
    if (Count > 0)
    {
        FVector2D ViewportSize;
        GEngine->GameViewport->GetViewportSize(ViewportSize);
        auto Width = CVPixelBufferGetWidth(pixelBuffer);
        auto Height = CVPixelBufferGetHeight(pixelBuffer);
        FVector2D Scale = ViewportSize / FVector2D(Width, Height);
        UE_LOG(LogTemp, Log, TEXT("PixelBufferSize %d, %d ViewportSize %s, Scale %s"), Width, Height, *ViewportSize.ToString(), *Scale.ToString());
        //NSArray* ColorList = [Pipeline readOutputListString:@"colors"];
        [BoxList enumerateObjectsUsingBlock:^(id object, NSUInteger idx, BOOL *stop) {
            NSString* MakeModel = [MakeModelList objectAtIndex: idx];
            NSString* Color = [ColorList objectAtIndex: idx];
            CGRect BBox = [object CGRectValue];
            NSLog(@"Make: %@ (len=%lu) Color %@ Box.x %f Box.y %f Box.w %f Box.h %f", MakeModel, MakeModel.length, Color, BBox.origin.x, BBox.origin.y, BBox.size.width, BBox.size.height);
            if (MakeModel.length > 1)
            {
                FCarModelRecognition Rec;
                Rec.MakeModel = FString(MakeModel);
                //UE_LOG(LogTemp, Log, TEXT("Car recognized: %s"), *Rec.MakeModel);
                float CropX = ((Scale.Y * Width)-ViewportSize.X)/2.0f;
                FVector2D Origin((BBox.origin.x/Width)*ViewportSize.X-CropX, (BBox.origin.y/Height)*ViewportSize.Y);
                FVector2D Size((BBox.size.width/Width)*ViewportSize.X, (BBox.size.height/Height)*ViewportSize.Y);
                Rec.Center = Origin + Size/2.0f;
                Rec.Extent = Size/2.0f;
                ResultPtr->Add(Rec);
            }
        }];
    }
}
#endif

void USighthoundRecognizerComponent::FireCarModelsRecognized(const TArray<FCarModelRecognition>& CarModelsRecognized)
{
    OnCarModelsRecognized.Broadcast(CarModelsRecognized);
}

// Sets default values for this component's properties
USighthoundRecognizerComponent::USighthoundRecognizerComponent()
{
	// Set this component to be initialized when the game starts, and to be ticked every frame.  You can turn these features
	// off to improve performance if you don't need them.
	PrimaryComponentTick.bCanEverTick = true;

	// ...
}


// Called when the game starts
void USighthoundRecognizerComponent::BeginPlay()
{
	Super::BeginPlay();

	// ...
	
}

void USighthoundRecognizerComponent::Apply(const TArray<FColor>& BgraPixels, int32 Width, int32 Height)
{
#if PLATFORM_IOS
    NSDictionary *options = nil;
    
    //CVPixelBufferCreateWithBytes(CFAllocatorRef allocator, size_t width, size_t height, OSType pixelFormatType, void *baseAddress, size_t bytesPerRow, CVPixelBufferReleaseBytesCallback releaseCallback, void *releaseRefCon, CFDictionaryRef pixelBufferAttributes, CVPixelBufferRef  _Nullable *pixelBufferOut);
    if (outPixelBuffer == nullptr)
    {
        CVReturn status =
        CVPixelBufferCreateWithBytes(
                                     kCFAllocatorDefault, size_t(Width), size_t(Height),
                                     kCVPixelFormatType_32BGRA, (void*)BgraPixels.GetData(), size_t(Width*sizeof(FColor)), nullptr, nullptr, (__bridge CFDictionaryRef)options,
                                     &outPixelBuffer);
        if (status != kCVReturnSuccess)
        {
            return;
        }
    }
    else
    {
        CVPixelBufferLockBaseAddress(outPixelBuffer, 0);
        memcpy(CVPixelBufferGetBaseAddress(outPixelBuffer), BgraPixels.GetData(), BgraPixels.Num());
        CVPixelBufferUnlockBaseAddress(outPixelBuffer, 0);
    
        static bool bSaved = false;
        if (!bSaved)
        {
            bSaved = true;
            CGImageRef destRef;
            if (VTCreateCGImageFromCVPixelBuffer(outPixelBuffer, NULL, &destRef) == noErr)
            {
                UIImage* Image = [[UIImage alloc] initWithCGImage:destRef];
                CGImageRelease(destRef);
                [[PHPhotoLibrary sharedPhotoLibrary] performChanges:^{
                    [PHAssetChangeRequest creationRequestForAssetFromImage:Image];
                } completionHandler:^(BOOL success, NSError *error) {
                    if (success) {
                        NSLog(@"Successfully copied image to photos folder");
                    }
                    else {
                        NSLog(@"Failed to copy imge to photos folder : %@",error);
                    }
                    CFRelease(Image);
                }];
            }
        }
    }
    TSharedPtr<TArray<FCarModelRecognition>, ESPMode::ThreadSafe> ResultPtr(new TArray<FCarModelRecognition>());
    RunRecognizer(outPixelBuffer, ResultPtr);
    FireCarModelsRecognized(*ResultPtr);
#endif
}


// Called every frame
void USighthoundRecognizerComponent::TickComponent(float DeltaTime, ELevelTick TickType, FActorComponentTickFunction* ThisTickFunction)
{
	Super::TickComponent(DeltaTime, TickType, ThisTickFunction);
    if (bRecognitionStarted && !bRecognizerBusy)
    {
#if PLATFORM_IOS
        // ...
        UAppleARKitSession* Session = UAppleARKitBlueprintFunctionLibrary::GetSession();
        if (Session != nullptr)
        {
            double ImageTimestamp;
            auto CapturedImage = Session->GetCapturedImage(ImageTimestamp);
            if (CapturedImage != nullptr && LastRecognizedTimestamp != ImageTimestamp)
            {
                LastRecognizedTimestamp = ImageTimestamp;
                //UE_LOG(LogTemp, Log, TEXT("New Timestamp %f"), ImageTimestamp);
                {
                    bRecognizerBusy = true;
                    {
                        auto Rgb = Yuv420ToRgb(CapturedImage);
                        //FFunctionGraphTask::CreateAndDispatchWhenReady([=]() {
                            if (Rgb != nullptr)
                            {
                                TSharedPtr<TArray<FCarModelRecognition>, ESPMode::ThreadSafe> ResultPtr(new TArray<FCarModelRecognition>());
                                RunRecognizer(Rgb, ResultPtr);
                                CFRelease(Rgb);
                                
                                //FFunctionGraphTask::CreateAndDispatchWhenReady([=]() {
                                    FireCarModelsRecognized(*ResultPtr);
                                //}, TStatId(), nullptr, ENamedThreads::GameThread);
                                
                                //UE_LOG(LogTemp, Log, TEXT("Released Image"));
                            }
                            bRecognizerBusy = false;
                        //}, TStatId(), nullptr, ENamedThreads::ActualRenderingThread);
                    }
                }
            }
        }
#endif
    }
}