#!/usr/bin/env python
# coding: utf-8

# # Relabeling a label image
# See also [discussion on image.sc](https://forum.image.sc/t/making-np-where-faster/58641/2).

# In[1]:


import numpy as np
import numpy
import cupy as cp
import cupy
import pyclesperanto_prototype as cle
import timeit
cle.select_device("RTX")


# In[2]:


sliceimage = (np.random.random((17, 2439, 2439)) * 4).astype(int)

def relabel_numpy(sliceimage):
    return numpy.take(numpy.array([0, 3, 3, 6, 4], numpy.uint8), sliceimage)

def relabel_cupy(sliceimage):
    return cupy.asnumpy(
        cupy.take(
            cupy.asarray([0, 3, 3, 6, 4]), 
            cupy.asarray(sliceimage)
        )
    )

def relabel_cle(sliceimage):
    return cle.replace_intensities(sliceimage, np.asarray([0, 3, 3, 6, 4]))


# In[3]:


get_ipython().run_line_magic('timeit', 'relabel_cupy(sliceimage)')


# In[4]:


get_ipython().run_line_magic('timeit', 'relabel_numpy(sliceimage)')


# In[5]:


get_ipython().run_line_magic('timeit', 'relabel_cle(sliceimage)')


# In[ ]: