Add takeSnapshot support to EngineView (#77)

* Android takeSnapshot native code compiling, not hooked up yet.

* Fix useRef declaration.

* Add EngineViewHooks for registration of callbacks, dummy promise resolution for now.

* Screenshots working on Android

* Remove some debugging changes.

* Clean up some spacing

* iOS screenshots mostly working.  Need to investigate hang when grabbing current drawable.

* Use UIView approach rather than MTKView approach, so we don't have a race condition on the texture being available.

* Add some comments and clean up EngineViewManager.mm

* Add toggle for whether or not snapshots are enabled.

* Change afterScreenUpdates to no to immediately grab the pixels rather than waiting for the next frame, fix capture quality value.

* Add version checks and error handling to Android engineview implementation, and implement the non-deprecated version of receiveCommand.

* Change engineViewHooks to a useState, change onInitialized, onSnapshot, snapshotDataReturnedHandler to useEffects or useCallbacks as appropriate.

* Apply suggestions from code review

Co-authored-by: Ryan Tremblay <[email protected]>

* Address PR Comments

* Commit suggestion that somehow didn't get applied.

* Fix broken incorrect return statement when overlapping snapshots are requested.

Co-authored-by: Ryan Tremblay <[email protected]>

Author: Alex-MSFT

Apps/Playground/App.tsx

@@ -6,14 +6,15 @@
  */
 
 import React, { useState, FunctionComponent, useEffect, useCallback } from 'react';
-import { SafeAreaView, StatusBar, Button, View, Text, ViewProps } from 'react-native';
+import { SafeAreaView, StatusBar, Button, View, Text, ViewProps, Image } from 'react-native';
 
-import { EngineView, useEngine } from '@babylonjs/react-native';
-import { Scene, Vector3, Mesh, ArcRotateCamera, Camera, PBRMetallicRoughnessMaterial, Color3, TargetCamera, WebXRSessionManager } from '@babylonjs/core';
+import { EngineView, useEngine, EngineViewCallbacks } from '@babylonjs/react-native';
+import { Scene, Vector3, Mesh, ArcRotateCamera, Camera, PBRMetallicRoughnessMaterial, Color3, TargetCamera, WebXRSessionManager, Engine } from '@babylonjs/core';
 import Slider from '@react-native-community/slider';
 
 const EngineScreen: FunctionComponent<ViewProps> = (props: ViewProps) => {
   const defaultScale = 1;
+  const enableSnapshots = false;
 
   const engine = useEngine();
   const [toggleView, setToggleView] = useState(false);
@@ -22,6 +23,8 @@ const EngineScreen: FunctionComponent<ViewProps> = (props: ViewProps) => {
   const [scene, setScene] = useState<Scene>();
   const [xrSession, setXrSession] = useState<WebXRSessionManager>();
   const [scale, setScale] = useState<number>(defaultScale);
+  const [snapshotData, setSnapshotData] = useState<string>();
+  const [engineViewCallbacks, setEngineViewCallbacks] = useState<EngineViewCallbacks>();
 
   useEffect(() => {
     if (engine) {
@@ -65,21 +68,37 @@ const EngineScreen: FunctionComponent<ViewProps> = (props: ViewProps) => {
           // TODO: Figure out why getFrontPosition stopped working
           //box.position = (scene.activeCamera as TargetCamera).getFrontPosition(2);
           const cameraRay = scene.activeCamera!.getForwardRay(1);
-          box.position = cameraRay.origin.add(cameraRay.direction.scale(cameraRay.length)); 
+          box.position = cameraRay.origin.add(cameraRay.direction.scale(cameraRay.length));
           box.rotate(Vector3.Up(), 3.14159);
         }
       }
     })();
   }, [box, scene, xrSession]);
 
+  const onInitialized = useCallback(async(engineViewCallbacks: EngineViewCallbacks) => {
+    setEngineViewCallbacks(engineViewCallbacks);
+  }, [engine]);
+
+  const onSnapshot = useCallback(async () => {
+    if (engineViewCallbacks) {
+      setSnapshotData("data:image/jpeg;base64," + await engineViewCallbacks.takeSnapshot());
+    }
+  }, [engineViewCallbacks]);
+
   return (
     <>
       <View style={props.style}>
         <Button title="Toggle EngineView" onPress={() => { setToggleView(!toggleView) }} />
         <Button title={ xrSession ? "Stop XR" : "Start XR"} onPress={onToggleXr} />
         { !toggleView &&
           <View style={{flex: 1}}>
-            <EngineView style={props.style} camera={camera} />
+            { enableSnapshots && 
+              <View style ={{flex: 1}}>
+                <Button title={"Take Snapshot"} onPress={onSnapshot}/>
+                <Image style={{flex: 1}} source={{uri: snapshotData }} />
+              </View>
+            }
+            <EngineView style={props.style} camera={camera} onInitialized={onInitialized} />
             <Slider style={{position: 'absolute', minHeight: 50, margin: 10, left: 0, right: 0, bottom: 0}} minimumValue={0.2} maximumValue={2} value={defaultScale} onValueChange={setScale} />
           </View>
         }

Modules/@babylonjs/react-native/EngineView.tsx

@@ -1,5 +1,5 @@
-import React, { FunctionComponent, Component, useEffect, useState } from 'react';
-import { requireNativeComponent, NativeModules, ViewProps, AppState, AppStateStatus, View, Text } from 'react-native';
+import React, { Component, FunctionComponent, SyntheticEvent, useCallback, useEffect, useState, useRef } from 'react';
+import { requireNativeComponent, NativeModules, ViewProps, AppState, AppStateStatus, View, Text, findNodeHandle, UIManager } from 'react-native';
 import { Camera } from '@babylonjs/core';
 import { IsEngineDisposed } from './EngineHelpers';
 import { BabylonModule } from './BabylonModule';
@@ -17,6 +17,7 @@ if (EngineViewManager && EngineViewManager.setJSThread && !isRemoteDebuggingEnab
 }
 
 interface NativeEngineViewProps extends ViewProps {
+    onSnapshotDataReturned: (event: SyntheticEvent) => void;
 }
 
 const NativeEngineView: {
@@ -27,11 +28,18 @@ const NativeEngineView: {
 export interface EngineViewProps extends ViewProps {
     camera?: Camera;
     displayFrameRate?: boolean;
+    onInitialized?: (view: EngineViewCallbacks) => void;
+}
+
+export interface EngineViewCallbacks {
+    takeSnapshot: () => Promise<string>;
 }
 
 export const EngineView: FunctionComponent<EngineViewProps> = (props: EngineViewProps) => {
     const [failedInitialization, setFailedInitialization] = useState(false);
     const [fps, setFps] = useState<number>();
+    const engineViewRef = useRef<Component<NativeEngineViewProps>>(null);
+    const snapshotPromise = useRef<{promise: Promise<string>, resolve: (data: string) => void}>();
 
     useEffect(() => {
         (async () => {
@@ -90,10 +98,50 @@ export const EngineView: FunctionComponent<EngineViewProps> = (props: EngineView
         setFps(undefined);
     }, [props.camera, props.displayFrameRate]);
 
+    // Call onInitialized if provided, and include the callback for takeSnapshot.
+    useEffect(() => {
+        if (props.onInitialized) {
+            props.onInitialized({
+                takeSnapshot: (): Promise<string> => {
+                    if (!snapshotPromise.current) {
+                        let resolveFunction: ((data: string) => void) | undefined;
+                        const promise = new Promise<string>((resolutionFunc) => {
+                            resolveFunction = resolutionFunc;
+                        });
+
+                        // Resolution functions should always be initialized.
+                        if (resolveFunction) {
+                            snapshotPromise.current = { promise: promise, resolve: resolveFunction };
+                        }
+                        else {
+                            throw "Resolution functions not initialized after snapshot promise creation.";
+                        }
+
+                        UIManager.dispatchViewManagerCommand(
+                            findNodeHandle(engineViewRef.current),
+                            "takeSnapshot",
+                            []);
+                    }
+
+                    return snapshotPromise.current.promise;
+                }
+            });
+        }
+    }, [props.onInitialized]);
+
+    // Handle snapshot data returned.
+    const snapshotDataReturnedHandler = useCallback((event: SyntheticEvent) => {
+        const { data } = event.nativeEvent;
+        if (snapshotPromise.current) {
+            snapshotPromise.current.resolve(data);
+            snapshotPromise.current = undefined;
+        }
+    }, []);
+
     if (!failedInitialization) {
         return (
             <View style={[props.style, {overflow: "hidden"}]}>
-                <NativeEngineView style={{flex: 1}} />
+                <NativeEngineView ref={engineViewRef} style={{flex: 1}} onSnapshotDataReturned={snapshotDataReturnedHandler} />
                 { fps && <Text style={{color: 'yellow', position: 'absolute', margin: 10, right: 0, top: 0}}>FPS: {Math.round(fps)}</Text> }
             </View>
         );
@@ -105,4 +153,4 @@ export const EngineView: FunctionComponent<EngineViewProps> = (props: EngineView
             </View>
         );
     }
-}
+}

Modules/@babylonjs/react-native/android/src/main/java/com/reactlibrary/EngineView.java

@@ -1,20 +1,33 @@
 package com.reactlibrary;
 
+import android.annotation.TargetApi;
+import android.graphics.Bitmap;
+import android.os.Build;
+import android.os.Handler;
+import android.os.HandlerThread;
+import android.util.Base64;
 import android.view.MotionEvent;
+import android.view.PixelCopy;
 import android.view.SurfaceHolder;
 import android.view.SurfaceView;
 import android.view.View;
 
 import com.facebook.react.bridge.ReactContext;
+import com.facebook.react.uimanager.UIManagerModule;
+import com.facebook.react.uimanager.events.EventDispatcher;
+
+import java.io.ByteArrayOutputStream;
 
 public final class EngineView extends SurfaceView implements SurfaceHolder.Callback, View.OnTouchListener {
     private final ReactContext reactContext;
+    private final EventDispatcher reactEventDispatcher;
 
     public EngineView(ReactContext reactContext) {
         super(reactContext);
         this.reactContext = reactContext;
         this.getHolder().addCallback(this);
         this.setOnTouchListener(this);
+        this.reactEventDispatcher = this.reactContext.getNativeModule(UIManagerModule.class).getEventDispatcher();
     }
 
     @Override
@@ -37,4 +50,41 @@ public boolean onTouch(View view, MotionEvent motionEvent) {
         BabylonNativeInterop.reportMotionEvent(this.reactContext, motionEvent);
         return true;
     }
+
+    @TargetApi(24)
+    public void takeSnapshot() {
+        // Only supported on API level 24 and up, return a blank image.
+        if (android.os.Build.VERSION.SDK_INT < Build.VERSION_CODES.N) {
+            SnapshotDataReturnedEvent snapshotEvent = new SnapshotDataReturnedEvent(this.getId(), "");
+            reactEventDispatcher.dispatchEvent(snapshotEvent);
+        }
+
+        // Create a bitmap that matches the width and height of the EngineView.
+        final Bitmap bitmap = Bitmap.createBitmap(
+                getWidth(),
+                getHeight(),
+                Bitmap.Config.ARGB_8888);
+
+        // Offload the snapshot worker to a helper thread.
+        final HandlerThread helperThread = new HandlerThread("ScreenCapture",-1);
+        helperThread.start();
+        final Handler helperThreadHandler = new Handler(helperThread.getLooper());
+
+        // Request the pixel copy.
+        PixelCopy.request(this, bitmap, (copyResult) ->  {
+            // If the pixel copy was a success then convert the image to a base 64 encoded jpeg and fire the event.
+            String encoded = "";
+            if (copyResult == PixelCopy.SUCCESS) {
+                ByteArrayOutputStream byteArrayStream = new ByteArrayOutputStream();
+                bitmap.compress(Bitmap.CompressFormat.JPEG, 100, byteArrayStream);
+                byte[] byteArray = byteArrayStream.toByteArray();
+                bitmap.recycle();
+                encoded = Base64.encodeToString(byteArray, Base64.DEFAULT);
+            }
+
+            SnapshotDataReturnedEvent snapshotEvent = new SnapshotDataReturnedEvent(this.getId(), encoded);
+            reactEventDispatcher.dispatchEvent(snapshotEvent);
+            helperThread.quitSafely();
+        }, helperThreadHandler);
+    }
 }

Modules/@babylonjs/react-native/android/src/main/java/com/reactlibrary/EngineViewManager.java

@@ -3,13 +3,16 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
+import com.facebook.react.bridge.ReadableArray;
 import com.facebook.react.common.MapBuilder;
 import com.facebook.react.uimanager.SimpleViewManager;
 import com.facebook.react.uimanager.ThemedReactContext;
 
 import java.util.Map;
 
 public final class EngineViewManager extends SimpleViewManager<EngineView> {
+    public static final int COMMAND_TAKE_SNAPSHOT = 0;
+    public static final String COMMAND_TAKE_SNAPSHOT_NAME = "takeSnapshot";
 
     @NonNull
     @Override
@@ -28,4 +31,33 @@ public void onDropViewInstance(@NonNull EngineView view) {
         super.onDropViewInstance(view);
         // TODO: Native view specific cleanup
     }
+
+    @Override
+    public Map<String,Integer> getCommandsMap() {
+        return MapBuilder.of(
+            COMMAND_TAKE_SNAPSHOT_NAME,
+            COMMAND_TAKE_SNAPSHOT
+        );
+    }
+
+    @Override
+    public void receiveCommand(final EngineView view, String commandId, ReadableArray args) {
+        // This will be called whenever a command is sent from react-native.
+        switch (commandId) {
+            case COMMAND_TAKE_SNAPSHOT_NAME:
+                view.takeSnapshot();
+                break;
+            default:
+                throw new IllegalArgumentException(
+                    String.format("Invalid command %s specified for EngineView.  Supported Commands: takeSnapshot", commandId));
+        }
+    }
+
+    @Override
+    public Map getExportedCustomDirectEventTypeConstants() {
+        return MapBuilder.builder()
+                .put(SnapshotDataReturnedEvent.EVENT_NAME,
+                    MapBuilder.of("registrationName", SnapshotDataReturnedEvent.EVENT_NAME))
+                .build();
+    }
 }

Modules/@babylonjs/react-native/android/src/main/java/com/reactlibrary/SnapshotDataReturnedEvent.java

@@ -0,0 +1,28 @@
+package com.reactlibrary;
+
+import com.facebook.react.bridge.Arguments;
+import com.facebook.react.bridge.WritableMap;
+import com.facebook.react.uimanager.events.Event;
+import com.facebook.react.uimanager.events.RCTEventEmitter;
+
+public class SnapshotDataReturnedEvent extends Event<SnapshotDataReturnedEvent> {
+    public static final String EVENT_NAME = "onSnapshotDataReturned";
+    public static final String DATA_NAME = "data";
+    private final WritableMap payload;
+
+    public SnapshotDataReturnedEvent(int viewId, String imageData) {
+        super(viewId);
+        this.payload = Arguments.createMap();
+        this.payload.putString(DATA_NAME, imageData);
+    }
+
+    @Override
+    public String getEventName() {
+        return EVENT_NAME;
+    }
+
+    @Override
+    public void dispatch(RCTEventEmitter rctEventEmitter) {
+        rctEventEmitter.receiveEvent(getViewTag(), getEventName(), this.payload);
+    }
+}

Modules/@babylonjs/react-native/ios/EngineViewManager.mm

@@ -1,12 +1,16 @@
 #import "BabylonNativeInterop.h"
 
 #import <React/RCTViewManager.h>
+#import <React/RCTUIManager.h>
 
 #import <Foundation/Foundation.h>
 #import <UIKit/UIKit.h>
 #import <MetalKit/MetalKit.h>
 
 @interface EngineView : MTKView
+
+@property (nonatomic, copy) RCTDirectEventBlock onSnapshotDataReturned;
+
 @end
 
 @implementation EngineView {
@@ -47,6 +51,28 @@ - (void)touchesCancelled:(NSSet<UITouch*>*)touches withEvent:(UIEvent*)event {
     [BabylonNativeInterop reportTouchEvent:touches withEvent:event];
 }
 
+- (void)takeSnapshot {
+    // We must take the screenshot on the main thread otherwise we might fail to get a valid handle on the view's image.
+    dispatch_async(dispatch_get_main_queue(), ^{
+        // Start the graphics context.
+        UIGraphicsBeginImageContextWithOptions(self.bounds.size, YES /* opaque */, 0.0f);
+        
+        // Draw the current state of the view into the graphics context.
+        [self drawViewHierarchyInRect:self.bounds afterScreenUpdates:NO];
+        
+        // Grab the image from the graphics context, and convert into a base64 encoded JPG.
+        UIImage* capturedImage = UIGraphicsGetImageFromCurrentImageContext();
+        UIGraphicsEndImageContext();
+        NSData* jpgData = UIImageJPEGRepresentation(capturedImage, 1.0f);
+        NSString* encodedData = [jpgData base64EncodedStringWithOptions:0];
+        
+        // Fire the onSnapshotDataReturned event if hooked up.
+        if (self.onSnapshotDataReturned != nil) {
+            self.onSnapshotDataReturned(@{ @"data":encodedData});
+        }
+    });
+}
+
 @end
 
 
@@ -59,6 +85,19 @@ @implementation EngineViewManager {
 
 RCT_EXPORT_MODULE(EngineViewManager)
 
+RCT_EXPORT_VIEW_PROPERTY(onSnapshotDataReturned, RCTDirectEventBlock)
+
+RCT_EXPORT_METHOD(takeSnapshot:(nonnull NSNumber*) reactTag) {
+    // Marshal the takeSnapshot call to the appropriate EngineView.
+    [self.bridge.uiManager addUIBlock:^(RCTUIManager *uiManager, NSDictionary<NSNumber*,UIView*>* viewRegistry) {
+        EngineView* view = (EngineView*)viewRegistry[reactTag];
+        if (!view || ![view isKindOfClass:[EngineView class]]) {
+            return;
+        }
+        [view takeSnapshot];
+    }];
+}
+
 RCT_EXPORT_BLOCKING_SYNCHRONOUS_METHOD(setJSThread) {
     runLoop = [NSRunLoop currentRunLoop];
     return nil;