ARTest-Recorder.cs

Recorder.cs

using System;
using System.IO;
using UnityEngine;
using UnityEngine.Rendering;
using UnityEngine.XR.ARFoundation;
using UnityEngine.XR.ARSubsystems;
using Unity.Collections.LowLevel.Unsafe;
using ARKitStream.Internal;

// Save AR Related data and camera images
namespace ARKitStream
{
    public sealed class Recorder : MonoBehaviour
    {
        [SerializeField] ARCameraManager cameraManager = null;
        internal event Action<ARKitRemotePacket> PacketTransformer;
        Texture2D mTexture;
        ARKitService service;
        XRCameraImage image;
        Texture2D Tex2D;
        byte[] bytes;
        String ext;
        recordController recordController;
        public string SavePath;

        void Start()
        {
            if (Application.isEditor)
            {
                Destroy(gameObject);
                return;
            }

            recordController = GameObject.Find("Record Button").GetComponent<recordController>();

            // Set event for cameraManager
            cameraManager.frameReceived += OnCameraFrameReceived;
            InitSubSenders();
        }

        void OnDestroy()
        {
            if (cameraManager != null)
            {
                cameraManager.frameReceived -= OnCameraFrameReceived;
            }
        }

        void OnValidate()
        {
            if (cameraManager == null)
            {
                cameraManager = FindObjectOfType<ARCameraManager>();
            }
        }

        // Event with getting ARCameraFrame for cameraManager
        unsafe void OnCameraFrameReceived(ARCameraFrameEventArgs args)
        {
            // Get AR data
            var packet = new ARKitRemotePacket()
            {
                cameraFrame = new ARKitRemotePacket.CameraFrameEvent()
                {
                    timestampNs = args.timestampNs.Value,
                    projectionMatrix = args.projectionMatrix.Value,
                    displayMatrix = args.displayMatrix.Value
                }
            };
            if (PacketTransformer != null)
            {
                PacketTransformer(packet);
            }

            byte[] data = packet.Serialize();
            if (recordController.record == true)
            {

                // Save AR data
                SafeCreateDirectory(Application.persistentDataPath + "/" + recordController.time);
                saveARtoFile(data);

                // Get camera image
                if (!cameraManager.TryGetLatestImage(out image))
                {
                    return;
                }

                var conversionParams = new XRCameraImageConversionParams
                {
                    inputRect = new RectInt(0, 0, image.width, image.height),
                    outputDimensions = new Vector2Int(image.width / 2, image.height / 2),
                    outputFormat = TextureFormat.RGBA32,
                    transformation = CameraImageTransformation.MirrorX
                };

                if (mTexture == null)
                {
                    mTexture = new Texture2D(conversionParams.outputDimensions.x,
                                             conversionParams.outputDimensions.y,
                                             conversionParams.outputFormat, false);
                }

                // Save camera image
                var buffer = mTexture.GetRawTextureData<byte>();
                image.Convert(conversionParams, new IntPtr(buffer.GetUnsafePtr()), buffer.Length);
                mTexture.Apply();
                WriteTextureToFile(mTexture, args.timestampNs.ToString(), ImageFormat.JPEG);
                buffer.Dispose();
                image.Dispose();

            }
        }

        // Save AR data
        void saveARtoFile(byte[] data)
        {
            SavePath = Application.persistentDataPath + "/" + recordController.time + "/saved-ardata.bytes";
            using (var fs = new FileStream(SavePath, FileMode.Append, FileAccess.Write, FileShare.Read))
            {
                fs.Write(data, 0, data.Length);
                for (int i = 0; i < 5; i++) fs.WriteByte(0x1e); // record separator
            }
        }

        // Save camera image
        public void WriteTextureToFile(Texture2D texture, string filename, ImageFormat imageFormat)
        {
            if (texture == null)
                throw new System.ArgumentNullException("texture");

            if (Tex2D == null)
            {
                Tex2D = new Texture2D(texture.height, texture.width);
            }

            for (int y = 0; y < texture.height; ++y)
            {
                for (int x = 0; x < texture.width; ++x)
                {
                    Tex2D.SetPixel(y, texture.width - 1 - x, texture.GetPixel(x, y));
                }
            }
            Tex2D.Apply();

            var path = Path.Combine(Application.persistentDataPath, recordController.time, "imgs", filename);
            SafeCreateDirectory(Path.Combine(Application.persistentDataPath, recordController.time, "imgs"));

            if (string.IsNullOrEmpty(path))
                throw new System.InvalidOperationException("No path specified");

            bytes = Tex2D.EncodeToJPG();
            ext = ".jpg";
            if (bytes.Length > 0)
            {
                File.WriteAllBytes(path + ext, bytes);
            }
        }

        public static DirectoryInfo SafeCreateDirectory(string path)
        {
            if (Directory.Exists(path))
            {
                return null;
            }
            return Directory.CreateDirectory(path);
        }

        void InitSubSenders()
        {
            TrackedPoseSender.TryCreate(this);
            ARKitFaceSender.TryCreate(this);
            ARKitOcclusionSender.TryCreate(this);
            ARKitPlaneSender.TryCreate(this);
            ARKitHumanBodySender.TryCreate(this);
            ARKitRaycastSender.TryCreate(this);
        }

    }
}