fix issue referencing and copying data

bioimage-io · Oct 3, 2024 · 24c5a81 · 24c5a81
1 parent c643162
commit 24c5a81
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 23 deletions.
diff --git a/src/main/java/io/bioimage/modelrunner/tensorflow/v2/api020/shm/ShmBuilder.java b/src/main/java/io/bioimage/modelrunner/tensorflow/v2/api020/shm/ShmBuilder.java
@@ -98,8 +98,11 @@ private static void buildFromTensorUByte(Tensor<TUint8> tensor, String memoryNam
 			throw new IllegalArgumentException("Model output tensor with shape " + Arrays.toString(arrayShape) 
 					+ " is too big. Max number of elements per ubyte output tensor supported: " + Integer.MAX_VALUE / 1);
         SharedMemoryArray shma = SharedMemoryArray.readOrCreate(memoryName, arrayShape, new UnsignedByteType(), false, true);
-        ByteBuffer buff1 = shma.getDataBufferNoHeader();
-        tensor.rawData().read(buff1.array(), 0, buff1.capacity());
+        ByteBuffer buff = shma.getDataBufferNoHeader();
+        byte[] flat = new byte[buff.capacity()];
+        ByteBuffer buff2 = ByteBuffer.wrap(flat);
+        tensor.rawData().read(flat, 0, buff.capacity());
+        buff = buff2;
         if (PlatformDetection.isWindows()) shma.close();
     }
 
@@ -111,8 +114,11 @@ private static void buildFromTensorInt(Tensor<TInt32> tensor, String memoryName)
 					+ " is too big. Max number of elements per int output tensor supported: " + Integer.MAX_VALUE / 4);
 
         SharedMemoryArray shma = SharedMemoryArray.readOrCreate(memoryName, arrayShape, new IntType(), false, true);
-        ByteBuffer buff1 = shma.getDataBufferNoHeader();
-        tensor.rawData().read(buff1.array(), 0, buff1.capacity());
+        ByteBuffer buff = shma.getDataBufferNoHeader();
+        byte[] flat = new byte[buff.capacity()];
+        ByteBuffer buff2 = ByteBuffer.wrap(flat);
+        tensor.rawData().read(flat, 0, buff.capacity());
+        buff = buff2;
         if (PlatformDetection.isWindows()) shma.close();
     }
 
@@ -124,8 +130,11 @@ private static void buildFromTensorFloat(Tensor<TFloat32> tensor, String memoryN
 					+ " is too big. Max number of elements per float output tensor supported: " + Integer.MAX_VALUE / 4);
 
         SharedMemoryArray shma = SharedMemoryArray.readOrCreate(memoryName, arrayShape, new FloatType(), false, true);
-        ByteBuffer buff1 = shma.getDataBufferNoHeader();
-        tensor.rawData().read(buff1.array(), 0, buff1.capacity());
+        ByteBuffer buff = shma.getDataBufferNoHeader();
+        byte[] flat = new byte[buff.capacity()];
+        ByteBuffer buff2 = ByteBuffer.wrap(flat);
+        tensor.rawData().read(flat, 0, buff.capacity());
+        buff = buff2;
         if (PlatformDetection.isWindows()) shma.close();
     }
 
@@ -137,8 +146,11 @@ private static void buildFromTensorDouble(Tensor<TFloat64> tensor, String memory
 					+ " is too big. Max number of elements per double output tensor supported: " + Integer.MAX_VALUE / 8);
 
         SharedMemoryArray shma = SharedMemoryArray.readOrCreate(memoryName, arrayShape, new DoubleType(), false, true);
-        ByteBuffer buff1 = shma.getDataBufferNoHeader();
-        tensor.rawData().read(buff1.array(), 0, buff1.capacity());
+        ByteBuffer buff = shma.getDataBufferNoHeader();
+        byte[] flat = new byte[buff.capacity()];
+        ByteBuffer buff2 = ByteBuffer.wrap(flat);
+        tensor.rawData().read(flat, 0, buff.capacity());
+        buff = buff2;
         if (PlatformDetection.isWindows()) shma.close();
     }
 
@@ -151,8 +163,11 @@ private static void buildFromTensorLong(Tensor<TInt64> tensor, String memoryName
 
 
         SharedMemoryArray shma = SharedMemoryArray.readOrCreate(memoryName, arrayShape, new LongType(), false, true);
-        ByteBuffer buff1 = shma.getDataBufferNoHeader();
-        tensor.rawData().read(buff1.array(), 0, buff1.capacity());
+        ByteBuffer buff = shma.getDataBufferNoHeader();
+        byte[] flat = new byte[buff.capacity()];
+        ByteBuffer buff2 = ByteBuffer.wrap(flat);
+        tensor.rawData().read(flat, 0, buff.capacity());
+        buff = buff2;
         if (PlatformDetection.isWindows()) shma.close();
     }
 }
diff --git a/src/main/java/io/bioimage/modelrunner/tensorflow/v2/api020/shm/TensorBuilder.java b/src/main/java/io/bioimage/modelrunner/tensorflow/v2/api020/shm/TensorBuilder.java
@@ -26,10 +26,6 @@
 import net.imglib2.util.Cast;
 
 import java.nio.ByteBuffer;
-import java.nio.DoubleBuffer;
-import java.nio.FloatBuffer;
-import java.nio.IntBuffer;
-import java.nio.LongBuffer;
 import java.util.Arrays;
 
 import org.tensorflow.Tensor;
@@ -102,7 +98,10 @@ private static Tensor<TUint8> buildUByte(SharedMemoryArray tensor)
 		if (!tensor.isNumpyFormat())
 			throw new IllegalArgumentException("Shared memory arrays must be saved in numpy format.");
 		ByteBuffer buff = tensor.getDataBufferNoHeader();
-		ByteDataBuffer dataBuffer = RawDataBufferFactory.create(buff.array(), false);
+		byte[] flat = new byte[buff.capacity()];
+		buff.get(flat);
+		buff.rewind();
+		ByteDataBuffer dataBuffer = RawDataBufferFactory.create(flat, false);
 		Tensor<TUint8> ndarray = Tensor.of(TUint8.DTYPE, Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}
@@ -117,8 +116,10 @@ private static Tensor<TInt32> buildInt(SharedMemoryArray tensor)
 		if (!tensor.isNumpyFormat())
 			throw new IllegalArgumentException("Shared memory arrays must be saved in numpy format.");
 		ByteBuffer buff = tensor.getDataBufferNoHeader();
-		IntBuffer intBuff = buff.asIntBuffer();
-		IntDataBuffer dataBuffer = RawDataBufferFactory.create(intBuff.array(), false);
+		int[] flat = new int[buff.capacity() / 4];
+		buff.asIntBuffer().get(flat);
+		buff.rewind();
+		IntDataBuffer dataBuffer = RawDataBufferFactory.create(flat, false);
 		Tensor<TInt32> ndarray = TInt32.tensorOf(Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}
@@ -133,8 +134,10 @@ private static Tensor<TInt64> buildLong(SharedMemoryArray tensor)
 		if (!tensor.isNumpyFormat())
 			throw new IllegalArgumentException("Shared memory arrays must be saved in numpy format.");
 		ByteBuffer buff = tensor.getDataBufferNoHeader();
-		LongBuffer longBuff = buff.asLongBuffer();
-		LongDataBuffer dataBuffer = RawDataBufferFactory.create(longBuff.array(), false);
+		long[] flat = new long[buff.capacity() / 8];
+		buff.asLongBuffer().get(flat);
+		buff.rewind();
+		LongDataBuffer dataBuffer = RawDataBufferFactory.create(flat, false);
 		Tensor<TInt64> ndarray = TInt64.tensorOf(Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}
@@ -149,8 +152,10 @@ private static Tensor<TFloat32> buildFloat(SharedMemoryArray tensor)
 		if (!tensor.isNumpyFormat())
 			throw new IllegalArgumentException("Shared memory arrays must be saved in numpy format.");
 		ByteBuffer buff = tensor.getDataBufferNoHeader();
-		FloatBuffer floatBuff = buff.asFloatBuffer();
-		FloatDataBuffer dataBuffer = RawDataBufferFactory.create(floatBuff.array(), false);
+		float[] flat = new float[buff.capacity() / 4];
+		buff.asFloatBuffer().get(flat);
+		buff.rewind();
+		FloatDataBuffer dataBuffer = RawDataBufferFactory.create(flat, false);
 		Tensor<TFloat32> ndarray = TFloat32.tensorOf(Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}
@@ -165,8 +170,10 @@ private static Tensor<TFloat64> buildDouble(SharedMemoryArray tensor)
 		if (!tensor.isNumpyFormat())
 			throw new IllegalArgumentException("Shared memory arrays must be saved in numpy format.");
 		ByteBuffer buff = tensor.getDataBufferNoHeader();
-		DoubleBuffer doubleBuff = buff.asDoubleBuffer();
-		DoubleDataBuffer dataBuffer = RawDataBufferFactory.create(doubleBuff.array(), false);
+		double[] flat = new double[buff.capacity() / 8];
+		buff.asDoubleBuffer().get(flat);
+		buff.rewind();
+		DoubleDataBuffer dataBuffer = RawDataBufferFactory.create(flat, false);
 		Tensor<TFloat64> ndarray = TFloat64.tensorOf(Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}