Merge pull request #5 from Clinical-Support-Systems/MLKitInfer-NewChanges

Refactor MLKitInfer

Author: kfrancis

src/Plugin.Maui.ML/Platforms/Android/MLKitInfer.cs

@@ -1,5 +1,6 @@
-#if ANDROID
+#if ANDROID
 using Java.Nio;
+using Java.Nio.Channels;
 using Microsoft.ML.OnnxRuntime.Tensors;
 using System.Reflection;
 using Xamarin.TensorFlow.Lite;
@@ -14,6 +15,9 @@ public sealed class MLKitInfer : IMLInfer, IDisposable
     private bool _disposed;
     private Interpreter? _interpreter;
     private byte[]? _modelBytes;
+#pragma warning disable CS0414 // Field is assigned intentionally to keep the mapped model buffer alive for the interpreter lifetime.
+    private ByteBuffer? _modelBuffer;
+#pragma warning restore CS0414
 
     public void Dispose()
     {
@@ -44,21 +48,55 @@ public async Task LoadModelAsync(Stream modelStream, CancellationToken cancellat
         Initialize(ms.ToArray());
     }
 
-    public async Task LoadModelFromAssetAsync(string assetName, CancellationToken cancellationToken = default)
+    public Task LoadModelFromAssetAsync(string assetName, CancellationToken cancellationToken = default)
     {
         if (string.IsNullOrWhiteSpace(assetName))
             throw new ArgumentException("Asset name cannot be null or empty", nameof(assetName));
-        try
-        {
-            await using var assetStream = Application.Context.Assets?.Open(assetName)
-                                          ?? throw new FileNotFoundException(
-                                              $"Asset '{assetName}' not found in Android assets.");
-            await LoadModelAsync(assetStream, cancellationToken).ConfigureAwait(false);
-        }
-        catch (Exception ex)
+
+        return Task.Run(() =>
         {
-            throw new InvalidOperationException($"Failed to load TFLite model asset '{assetName}': {ex.Message}", ex);
-        }
+            cancellationToken.ThrowIfCancellationRequested();
+
+            try
+            {
+                var assets = Application.Context.Assets
+                             ?? throw new InvalidOperationException("Android AssetManager is unavailable.");
+
+                try
+                {
+                    using var afd = assets.OpenFd(assetName)
+                                    ?? throw new FileNotFoundException(
+                                        $"Asset '{assetName}' not found in Android assets.");
+
+                    using var inputStream = new Java.IO.FileInputStream(afd.FileDescriptor);
+
+                    var channel = inputStream.Channel
+                                  ?? throw new InvalidOperationException("Could not get FileChannel for asset.");
+
+                    var mappedBuffer = channel.Map(
+                        FileChannel.MapMode.ReadOnly!,
+                        afd.StartOffset,
+                        afd.DeclaredLength);
+
+                    Initialize(mappedBuffer);
+                }
+                catch (Exception ex) when (ex is not OperationCanceledException)
+                {
+                    cancellationToken.ThrowIfCancellationRequested();
+
+                    using var assetStream = assets.Open(assetName)
+                        ?? throw new FileNotFoundException($"Asset '{assetName}' not found in Android assets.", ex);
+                    using var ms = new System.IO.MemoryStream();
+                    assetStream.CopyTo(ms);
+                    Initialize(ms.ToArray());
+                }
+            }
+            catch (Exception ex) when (ex is not OperationCanceledException)
+            {
+                throw new InvalidOperationException(
+                    $"Failed to load TFLite model asset '{assetName}': {ex.Message}", ex);
+            }
+        }, cancellationToken);
     }
 
     public Task<Dictionary<string, Tensor<float>>> RunInferenceAsync(Dictionary<string, Tensor<float>> inputs,
@@ -126,20 +164,71 @@ public Task<Dictionary<string, Tensor<float>>> RunInferenceAsync(Dictionary<stri
     }
 
     public Task<Dictionary<string, Tensor<float>>> RunInferenceLongInputsAsync(Dictionary<string, Tensor<long>> inputs,
-        CancellationToken cancellationToken = default)
+    CancellationToken cancellationToken = default)
     {
+        if (!IsModelLoaded) throw new InvalidOperationException("No TFLite model loaded. Call LoadModelAsync first.");
         if (inputs == null || inputs.Count == 0)
             throw new ArgumentException("Inputs cannot be null or empty", nameof(inputs));
-        var floatInputs = new Dictionary<string, Tensor<float>>();
-        foreach (var (k, v) in inputs)
+
+        return Task.Run(() =>
         {
-            var cast = new DenseTensor<float>(v.Dimensions.ToArray());
-            var arr = v.ToArray();
-            for (var i = 0; i < arr.Length; i++) cast.Buffer.Span[i] = arr[i];
-            floatInputs[k] = cast;
-        }
+            lock (_sync)
+            {
+                if (_interpreter == null) throw new InvalidOperationException("Interpreter disposed.");
+
+                var first = inputs.First();
+                var inputTensor = first.Value;
+                var inputIndex = 0; // assume single input
+
+                var flat = inputTensor.ToArray();
+                var inputShape = inputTensor.Dimensions.ToArray();
+
+                try
+                {
+                    _interpreter.ResizeInput(inputIndex, inputShape);
+                }
+                catch
+                {
+                    /* ignore if immutable */
+                }
+
+                _interpreter.AllocateTensors();
+
+                var nativeOrder = ByteOrder.NativeOrder() ?? ByteOrder.BigEndian;
+
+                // INT64 input buffer
+                var inputData = ByteBuffer.AllocateDirect(flat.Length * 8).Order(nativeOrder!);
+                inputData.AsLongBuffer()!.Put(flat);
+                inputData.Rewind();
+
+                var outputTensor = _interpreter.GetOutputTensor(0);
+                if (outputTensor == null)
+                    throw new InvalidOperationException("Failed to retrieve output tensor.");
+
+                var oshape = outputTensor.Shape();
+                if (oshape == null)
+                    throw new InvalidOperationException("Failed to retrieve output shape.");
+
+                var outCount = 1;
+                foreach (var d in oshape) outCount *= d;
 
-        return RunInferenceAsync(floatInputs, cancellationToken);
+                // FLOAT32 output buffer
+                var outputData = ByteBuffer.AllocateDirect(outCount * 4).Order(nativeOrder!);
+
+                _interpreter.Run(inputData, outputData);
+
+                outputData.Rewind();
+
+                var outputArray = new float[outCount];
+                outputData.AsFloatBuffer()!.Get(outputArray);
+
+                var dense = new DenseTensor<float>(oshape.Select(i => i).ToArray());
+                var span = dense.Buffer.Span;
+                for (var i = 0; i < outputArray.Length; i++) span[i] = outputArray[i];
+
+                return new Dictionary<string, Tensor<float>> { ["output0"] = dense };
+            }
+        }, cancellationToken);
     }
 
     public Dictionary<string, MLNodeMetadata> GetInputMetadata()
@@ -174,6 +263,7 @@ public void UnloadModel()
             _interpreter?.Dispose();
             _interpreter = null;
             _modelBytes = null;
+            _modelBuffer = null;
         }
     }
 
@@ -193,6 +283,23 @@ private void Initialize(byte[] bytes)
             _interpreter.AllocateTensors();
         }
     }
+    private void Initialize(ByteBuffer modelBuffer)
+    {
+        lock (_sync)
+        {
+            UnloadModel();
+
+            _modelBytes = null;
+            _modelBuffer = modelBuffer;
+
+            _interpreter = new Interpreter(
+                modelBuffer,
+                (Interpreter.Options?)new Interpreter.Options().SetNumThreads(
+                    Math.Max(1, Environment.ProcessorCount - 1)));
+
+            _interpreter.AllocateTensors();
+        }
+    }
 
     private static Dictionary<string, MLNodeMetadata> GetTensorMetadata(Interpreter interpreter, bool isInput)
     {