Show dask array in BigDataViewer

dask_example.py

import dask.array
import dask.array.image
import dask.array as da
import numpy as np
import os
import pathlib
import skimage.io
import subprocess
import sys
import tempfile
import time
import threading

tmp_dir = tempfile.mkdtemp()

import jnius_config
jnius_config.add_classpath( tmp_dir )

import imglyb
import imglyb.types
import imglyb.util
from jnius import PythonJavaClass, java_method, autoclass, JavaException, cast

IntUnsafe  = autoclass('net.imglib2.img.basictypelongaccess.unsafe.IntUnsafe')
Cell       = autoclass('net.imglib2.img.cell.Cell')
System     = autoclass('java.lang.System')

make_cell_grid_code = """
import java.util.Arrays;

import net.imglib2.img.basictypelongaccess.unsafe.IntUnsafe;
import net.imglib2.img.NativeImg;
import net.imglib2.img.basictypeaccess.IntAccess;
import net.imglib2.img.cell.Cell;
import net.imglib2.img.cell.CellGrid;
import net.imglib2.img.cell.LazyCellImg;
import net.imglib2.img.cell.LazyCellImg.Get;
import net.imglib2.type.NativeType;
import net.imglib2.type.numeric.integer.UnsignedIntType;

public class Helpers {

    public static CellGrid makeGrid( long[] dims, int[] cellSize ) {
        final CellGrid grid =  new CellGrid( dims, cellSize );
return grid;
    }

    public static < T extends NativeType< T >, A > LazyCellImg< T, A > lazyCellImg( final CellGrid grid, final T type, final Get< Cell< A > > get )
    {
        return new LazyCellImg<>( grid, type, get );
    }

    public static long[] cellMin( final CellGrid grid, long index )
    {
        long[] min = new long[ grid.numDimensions() ];
        grid.getCellGridPositionFlat( index, min );
        Arrays.setAll( min, d -> min[ d ] * grid.cellDimension( d ) );
//        for ( int d = 0; d < min.length; ++d )
//            min[ d ] *= grid.cellDimension( d );
        return min;
    }

    public static UnsignedIntType makeLinkedType( final NativeImg< ?, ? extends IntAccess > img )
    {
        return new UnsignedIntType( img );
    }


}
"""

fp = pathlib.Path( tmp_dir ) / 'Helpers.java'
print( tmp_dir )
with open( fp, 'w' ) as f:
    f.write( make_cell_grid_code )

javac = pathlib.Path( os.environ[ 'JAVA_HOME' ] ) / 'bin' / 'javac'
proc = subprocess.run(
    [ javac, '-cp', jnius_config.split_char.join( jnius_config.get_classpath() ), fp ],
    stdout=subprocess.PIPE,
    stderr=subprocess.PIPE)
if proc.returncode != 0:
    print("Failed!")
    print ( proc.stderr )

Helpers = autoclass('Helpers')

class FunctionAsRunnable(PythonJavaClass):
    __javainterfaces__ = ['java/lang/Runnable']

    def __init__(self, func):
        super(FunctionAsRunnable, self).__init__()
        self.func = func

    @java_method('()V', name='run')
    def run(self):
        self.func()

class CacheLoaderFromDaskArray(PythonJavaClass):
    __javainterfaces__ = ['net/imglib2/cache/CacheLoader']

    def __init__(self, dask_array, cell_grid):
        super(CacheLoaderFromDaskArray, self).__init__()
        self.dask_array = dask_array
        self.slices     = da.core.slices_from_chunks(self.dask_array.chunks)
        self.cell_grid  = cell_grid
        self.cache      = {}
        self.cache_lock = threading.RLock()

    @java_method('(Ljava/lang/Object;)Ljava/lang/Object;', name='get')
    def get(self, index):
        chunk      = np.ascontiguousarray(self.dask_array[self.slices[index]].compute())
        refGuard   = imglyb.util.ReferenceGuard(chunk)
        address    = chunk.ctypes.data
        with self.cache_lock:
            self.cache[address] = chunk
        try:
            pos = Helpers.cellMin(self.cell_grid, index)
            access = IntUnsafe(address, None)
            cell = Cell(chunk.shape[::-1], Helpers.cellMin(self.cell_grid, index), access)
        except JavaException as e:
            print('Name    --', e.classname)
            print('Message --', e.innermessage)
            print('Stack trace --', e.stacktrace)
            raise e
        print("Adding index", index, "at address", address, "-- cache has", len(self.cache), "entries")
        return cell

class CacheRemoverAddressDictionary(PythonJavaClass):
    __javainterfaces__ = ['net/imglib2/cache/CacheRemover']

    def __init__(self, address_dictionary, lock):
        super(CacheRemoverAddressDictionary, self).__init__()
        self.address_dictionary = address_dictionary
        self.lock               = lock

    @java_method('(Ljava/lang/Object;Ljava/lang/Object;)V', name='onRemoval')
    def onRemoval(self, key, value):
        address = value.getData().getAddres()
        with self.lock:
            del self.address_dictionary[address]
        print("Removed key", key, "at address", address, "-- cache has", len(self.address_dictionary), "entries")


LazyCellImg     = autoclass('net.imglib2.img.cell.LazyCellImg')
CellGrid        = autoclass('net.imglib2.img.cell.CellGrid')
ARGBType        = autoclass('net.imglib2.type.numeric.ARGBType')
UnsignedIntType = autoclass('net.imglib2.type.numeric.integer.UnsignedIntType')

# chunks = [2, 3, 4]
# shape  = [4, 6, 8]


data   = skimage.io.imread("https://github.com/hanslovsky/imglyb-learnathon/raw/master/notebooks/basics/data/emdata.jpg")
chunks = tuple( s // 5 for s in data.shape )


array = da.from_array(data, chunks=chunks)
array = array.astype(np.uint32).rechunk(chunks)

shape = array.shape

print(shape)
print(chunks)

slices = da.core.slices_from_chunks(array.chunks)

grid = Helpers.makeGrid(shape[::-1], chunks[::-1])
print("slices")

loader                             = CacheLoaderFromDaskArray(array, grid)
remover                            = CacheRemoverAddressDictionary(loader.cache, loader.cache_lock)
SoftRefLoaderCache                 = autoclass('net.imglib2.cache.ref.SoftRefLoaderCache')
CachedCellImg                      = autoclass('net.imglib2.cache.img.CachedCellImg')
GuardedStrongRefLoaderCache        = autoclass('net.imglib2.cache.ref.GuardedStrongRefLoaderCache')
GuardedStrongRefLoaderRemoverCache = autoclass('net.imglib2.cache.ref.GuardedStrongRefLoaderRemoverCache')
t                                  = UnsignedIntType()
cache                              = GuardedStrongRefLoaderRemoverCache(5).withLoader(loader).withRemover(remover) # SoftRefLoaderCache().withLoader(loader)
img                                = CachedCellImg(grid, t.getEntitiesPerPixel(), cache, IntUnsafe(0))
linkedType                         = Helpers.makeLinkedType(img)
img.setLinkedType(linkedType)
try:
    img.randomAccess().get()
except JavaException as e:
    print(e.stacktrace)
    raise e

BdvOptions = autoclass('bdv.util.BdvOptions')

print("Going to open bdv window")
bdv = imglyb.util.BdvFunctions.show(img, 'random', BdvOptions.options().is2D())
bdv.getBdvHandle().getSetupAssignments().getMinMaxGroups().get(0).setRange(0, 255)
print("Showing bdv")
vp = bdv.getBdvHandle().getViewerPanel()

# Keep Python running until user closes Bdv window
check = autoclass( 'net.imglib2.python.BdvWindowClosedCheck' )()
frame = cast( 'javax.swing.JFrame', autoclass( 'javax.swing.SwingUtilities' ).getWindowAncestor( vp ) )
frame.addWindowListener( check )

def sleeper():
    while check.isOpen():
        time.sleep( 0.5 )
        # print([k for (k,v) in loader.cache.items()])

t = threading.Thread( target=sleeper )
t.start()
t.join()

dask_example_copy_access.py

import ctypes
import dask.array
import dask.array.image
import dask.array as da
import numpy as np
import os
import pathlib
import skimage.io
import subprocess
import sys
import tempfile
import time
import threading

_decref = ctypes.pythonapi.Py_DecRef
_decref.argtypes = [ctypes.py_object]
_decref.restype = None

_incref = ctypes.pythonapi.Py_IncRef
_incref.argtypes = [ctypes.py_object]
_incref.restype = None

tmp_dir = tempfile.mkdtemp()

import jnius_config
jnius_config.add_classpath( tmp_dir )

import imglyb
import imglyb.types
import imglyb.util
from jnius import PythonJavaClass, java_method, autoclass, JavaException, cast

IntUnsafe        = autoclass('net.imglib2.img.basictypelongaccess.unsafe.IntUnsafe')
Cell             = autoclass('net.imglib2.img.cell.Cell')
System           = autoclass('java.lang.System')
Executors        = autoclass('java.util.concurrent.Executors')
VolatileIntArray = autoclass('net.imglib2.img.basictypeaccess.volatiles.array.VolatileIntArray')


make_cell_grid_code = """
import java.lang.Runnable;

import java.util.Arrays;
import java.util.concurrent.ExecutorService;

import net.imglib2.img.basictypelongaccess.unsafe.IntUnsafe;
import net.imglib2.img.NativeImg;
import net.imglib2.img.basictypeaccess.IntAccess;
import net.imglib2.img.cell.Cell;
import net.imglib2.img.cell.CellGrid;
import net.imglib2.img.cell.LazyCellImg;
import net.imglib2.img.cell.LazyCellImg.Get;
import net.imglib2.type.NativeType;
import net.imglib2.type.numeric.integer.UnsignedIntType;

public class Helpers {

    public static CellGrid makeGrid( long[] dims, int[] cellSize ) {
        final CellGrid grid =  new CellGrid( dims, cellSize );
return grid;
    }

    public static < T extends NativeType< T >, A > LazyCellImg< T, A > lazyCellImg( final CellGrid grid, final T type, final Get< Cell< A > > get )
    {
        return new LazyCellImg<>( grid, type, get );
    }

    public static long[] cellMin( final CellGrid grid, long index )
    {
        long[] min = new long[ grid.numDimensions() ];
        grid.getCellGridPositionFlat( index, min );
        Arrays.setAll( min, d -> min[ d ] * grid.cellDimension( d ) );
//        for ( int d = 0; d < min.length; ++d )
//            min[ d ] *= grid.cellDimension( d );
        return min;
    }

    public static UnsignedIntType makeLinkedType( final NativeImg< ?, ? extends IntAccess > img )
    {
        return new UnsignedIntType( img );
    }

    public static class OnFinalize
    {

        private final Runnable onFinalize;

        public OnFinalize( Runnable onFinalize )
        {
            this.onFinalize = onFinalize;
        }

        @Override
        public void finalize()
        {
            onFinalize.run();
        }
    }

    public static OnFinalize onFinalize( Runnable onFinalize )
    {
        return new OnFinalize( onFinalize );
    }

    public static OnFinalize onFinalize( Runnable onFinalize, ExecutorService es )
    {
        Runnable r = () -> es.submit( onFinalize );
        return onFinalize( r );
    }

    public static Runnable test( Runnable r )
    {
        return r::run;
    }

    public static void copyAccess( IntAccess source, IntAccess target, int size )
    {
        for ( int index = 0; index < size; ++index )
        {
           target.setValue( index, source.getValue( index ) ); 
        }
    }

}
"""

fp = pathlib.Path( tmp_dir ) / 'Helpers.java'
print( tmp_dir )
with open( fp, 'w' ) as f:
    f.write( make_cell_grid_code )

javac = pathlib.Path( os.environ[ 'JAVA_HOME' ] ) / 'bin' / 'javac'
proc = subprocess.run(
    [ javac, '-cp', jnius_config.split_char.join( jnius_config.get_classpath() ), fp ],
    stdout=subprocess.PIPE,
    stderr=subprocess.PIPE)
if proc.returncode != 0:
    print("Failed!")
    print ( proc.stderr )

Helpers = autoclass('Helpers')
print(type(Helpers))

class FunctionAsRunnable(PythonJavaClass):
    __javainterfaces__ = ['java/lang/Runnable']

    def __init__(self, func):
        super(FunctionAsRunnable, self).__init__()
        self.func = func

    @java_method('()V', name='run')
    def run(self):
        print("RUNNING OR WHAT?")
        self.func()

# Helpers.test( FunctionAsRunnable( lambda : print("OK TEST"))).run()

class CacheLoaderFromDaskArray(PythonJavaClass):
    __javainterfaces__ = ['net/imglib2/cache/CacheLoader']

    def __init__(self, dask_array, cell_grid):
        super(CacheLoaderFromDaskArray, self).__init__()
        self.dask_array = dask_array
        self.slices     = da.core.slices_from_chunks(self.dask_array.chunks)
        self.cell_grid  = cell_grid

    @java_method('(Ljava/lang/Object;)Ljava/lang/Object;', name='get')
    def get(self, index):
        chunk      = self.dask_array[self.slices[index]].compute()
        # np.ascontiguousarray(self.dask_array[self.slices[index]].compute())
        refGuard   = imglyb.util.ReferenceGuard(chunk)
        address    = chunk.ctypes.data

        try:
            pos = Helpers.cellMin(self.cell_grid, index)
            source = IntUnsafe(address, None)
            target = VolatileIntArray(chunk.size, True)
            Helpers.copyAccess(source, target, chunk.size)
            cell = Cell(chunk.shape[::-1], Helpers.cellMin(self.cell_grid, index), target)
        except JavaException as e:
            print('Name    --', e.classname)
            print('Message --', e.innermessage)
            print('Stack trace --', e.stacktrace)
            raise e
        return cell

LazyCellImg     = autoclass('net.imglib2.img.cell.LazyCellImg')
CellGrid        = autoclass('net.imglib2.img.cell.CellGrid')
ARGBType        = autoclass('net.imglib2.type.numeric.ARGBType')
UnsignedIntType = autoclass('net.imglib2.type.numeric.integer.UnsignedIntType')

# chunks = [2, 3, 4]
# shape  = [4, 6, 8]


data   = skimage.io.imread("https://github.com/hanslovsky/imglyb-learnathon/raw/master/notebooks/basics/data/emdata.jpg")
chunks = tuple( s // 5 for s in data.shape )


array = da.from_array(data, chunks=chunks)
array = array.astype(np.uint32).rechunk(chunks)

shape = array.shape

print(shape)
print(chunks)

slices = da.core.slices_from_chunks(array.chunks)

grid = Helpers.makeGrid(shape[::-1], chunks[::-1])
print("slices")

loader                             = CacheLoaderFromDaskArray(array, grid)
SoftRefLoaderCache                 = autoclass('net.imglib2.cache.ref.SoftRefLoaderCache')
CachedCellImg                      = autoclass('net.imglib2.cache.img.CachedCellImg')
t                                  = UnsignedIntType()
cache                              = SoftRefLoaderCache().withLoader(loader)
img                                = CachedCellImg(grid, t.getEntitiesPerPixel(), cache, IntUnsafe(0))
linkedType                         = Helpers.makeLinkedType(img)
img.setLinkedType(linkedType)
try:
    img.randomAccess().get()
except JavaException as e:
    print(e.stacktrace)
    raise e

BdvOptions = autoclass('bdv.util.BdvOptions')

print("Going to open bdv window")
bdv = imglyb.util.BdvFunctions.show(img, 'random', BdvOptions.options().is2D())
bdv.getBdvHandle().getSetupAssignments().getMinMaxGroups().get(0).setRange(0, 255)
print("Showing bdv")
vp = bdv.getBdvHandle().getViewerPanel()

# Keep Python running until user closes Bdv window
check = autoclass( 'net.imglib2.python.BdvWindowClosedCheck' )()
frame = cast( 'javax.swing.JFrame', autoclass( 'javax.swing.SwingUtilities' ).getWindowAncestor( vp ) )
frame.addWindowListener( check )

def sleeper():
    while check.isOpen():
        time.sleep( 0.5 )

t = threading.Thread( target=sleeper )
t.start()
t.join()

no-ref-count-increase.py

import sys

from jnius import autoclass, PythonJavaClass, java_method

class OneSupplier( PythonJavaClass ):
        __javainterfaces__ = [ 'java.util.function.IntSupplier' ]
        def __init__( self, *args, **kwargs ):
                super( OneSupplier, self ).__init__()

        @java_method('()I')
        def getAsInt(self):
            return 1

ArrayList = autoclass('java.util.ArrayList')
al = ArrayList()
os = OneSupplier()
l  = []
print("Ref count before", sys.getrefcount(os))
al.add(os)
print("Ref count after", sys.getrefcount(os))
l.append(os)
print("Ref count after", sys.getrefcount(os))