Writing an async REPL - Part 1

This is a first part in a series of blog post which explain how I implemented the ability to await code at the top level scope in the IPython REPL. Don't expect the second part soon, or bother me for it. I know I shoudl write it, but time is a rarte luxury.

It is an interesting adventure into how Python code get executed, and I must admit it changed quite a bit how I understand python code now days and made me even more excited about async/await in Python.

It should also dive quite a bit in the internals of Python/CPython if you ever are interested in what some of these things are.

In [1]:
# we cheat and deactivate the new IPython feature to match Python repl behavior
%autoawait False

Async or not async, that is the question

You might now have noticed it, but since Python 3.5 the following is valid Python syntax:

In [2]:
async def a_function():
    async with contextmanager() as f:
        result = await f.get('stuff')
        return result

So you've been curious and read a lot about asyncio, and may have come across a few new libraries like aiohttp and all hte aio-libs, heard about sans-io, read complaints and we can take differents approaches, and maybe even maybe do better. You vaguely understand the concept of loops and futures, the term coroutine is still unclear. So you decide to poke around yourself in the REPL.

In [3]:
import aiohttp
In [4]:
print(aiohttp.__version__)
coro_req = aiohttp.get('https://api.github.com')
coro_req
1.3.5
Out[4]:
<aiohttp.client._DetachedRequestContextManager at 0x1045289d8>
In [5]:
import asyncio
res = asyncio.get_event_loop().run_until_complete(coro_req)
In [6]:
res
Out[6]:
<ClientResponse(https://api.github.com) [200 OK]>
<CIMultiDictProxy('Server': 'GitHub.com', 'Date': 'Thu, 06 Apr 2017 19:49:20 GMT', 'Content-Type': 'application/json; charset=utf-8', 'Transfer-Encoding': 'chunked', 'Status': '200 OK', 'X-Ratelimit-Limit': '60', 'X-Ratelimit-Remaining': '50', 'X-Ratelimit-Reset': '1491508909', 'Cache-Control': 'public, max-age=60, s-maxage=60', 'Vary': 'Accept', 'Etag': 'W/"7dc470913f1fe9bb6c7355b50a0737bc"', 'X-Github-Media-Type': 'github.v3; format=json', 'Access-Control-Expose-Headers': 'ETag, Link, X-GitHub-OTP, X-RateLimit-Limit, X-RateLimit-Remaining, X-RateLimit-Reset, X-OAuth-Scopes, X-Accepted-OAuth-Scopes, X-Poll-Interval', 'Access-Control-Allow-Origin': '*', 'Content-Security-Policy': "default-src 'none'", 'Strict-Transport-Security': 'max-age=31536000; includeSubdomains; preload', 'X-Content-Type-Options': 'nosniff', 'X-Frame-Options': 'deny', 'X-Xss-Protection': '1; mode=block', 'Vary': 'Accept-Encoding', 'X-Served-By': 'a51acaae89a7607fd7ee967627be18e4', 'Content-Encoding': 'gzip', 'X-Github-Request-Id': '8182:3911:C50FFE:EF0636:58E69BC0')>
In [7]:
res.json()
Out[7]:
<generator object ClientResponse.json at 0x1052cd9e8>
In [8]:
json = asyncio.get_event_loop().run_until_complete(res.json())
json
Out[8]:
{'authorizations_url': 'https://api.github.com/authorizations',
 'code_search_url': 'https://api.github.com/search/code?q={query}{&page,per_page,sort,order}',
 'commit_search_url': 'https://api.github.com/search/commits?q={query}{&page,per_page,sort,order}',
 'current_user_authorizations_html_url': 'https://github.com/settings/connections/applications{/client_id}',
 'current_user_repositories_url': 'https://api.github.com/user/repos{?type,page,per_page,sort}',
 'current_user_url': 'https://api.github.com/user',
 'emails_url': 'https://api.github.com/user/emails',
 'emojis_url': 'https://api.github.com/emojis',
 'events_url': 'https://api.github.com/events',
 'feeds_url': 'https://api.github.com/feeds',
 'followers_url': 'https://api.github.com/user/followers',
 'following_url': 'https://api.github.com/user/following{/target}',
 'gists_url': 'https://api.github.com/gists{/gist_id}',
 'hub_url': 'https://api.github.com/hub',
 'issue_search_url': 'https://api.github.com/search/issues?q={query}{&page,per_page,sort,order}',
 'issues_url': 'https://api.github.com/issues',
 'keys_url': 'https://api.github.com/user/keys',
 'notifications_url': 'https://api.github.com/notifications',
 'organization_repositories_url': 'https://api.github.com/orgs/{org}/repos{?type,page,per_page,sort}',
 'organization_url': 'https://api.github.com/orgs/{org}',
 'public_gists_url': 'https://api.github.com/gists/public',
 'rate_limit_url': 'https://api.github.com/rate_limit',
 'repository_search_url': 'https://api.github.com/search/repositories?q={query}{&page,per_page,sort,order}',
 'repository_url': 'https://api.github.com/repos/{owner}/{repo}',
 'starred_gists_url': 'https://api.github.com/gists/starred',
 'starred_url': 'https://api.github.com/user/starred{/owner}{/repo}',
 'team_url': 'https://api.github.com/teams',
 'user_organizations_url': 'https://api.github.com/user/orgs',
 'user_repositories_url': 'https://api.github.com/users/{user}/repos{?type,page,per_page,sort}',
 'user_search_url': 'https://api.github.com/search/users?q={query}{&page,per_page,sort,order}',
 'user_url': 'https://api.github.com/users/{user}'}

It's a bit painful to pass everything to run_until_complete, you know how to write async-def function and pass this to an event loop:

In [9]:
loop = asyncio.get_event_loop()
run = loop.run_until_complete
url = 'https://api.github.com/rate_limit'

async def get_json(url):
    res = await aiohttp.get(url)
    return await res.json()

run(get_json(url))
Out[9]:
{'rate': {'limit': 60, 'remaining': 50, 'reset': 1491508909},
 'resources': {'core': {'limit': 60, 'remaining': 50, 'reset': 1491508909},
  'graphql': {'limit': 0, 'remaining': 0, 'reset': 1491511760},
  'search': {'limit': 10, 'remaining': 10, 'reset': 1491508220}}}

Good ! And the you wonder, why do I have to wrap thing ina function, if I have a default loop isn't it obvious what where I want to run my code ? Can't I await things directly ? So you try:

In [10]:
await aiohttp.get(url)
  File "<ipython-input-10-055eb13ed07d>", line 1
    await aiohttp.get(url)
                ^
SyntaxError: invalid syntax

What ? Oh that's right there is no way in Pyton to set a default loop... but a SyntaxError ? Well, that's annoying.

Outsmart Python

Hopefully you (in this case me), are in control of the REPL. You can bend it to your will. Sure you can do some things. First you try to remember how a REPL works:

In [11]:
mycode = """
a = 1
print('hey')
"""
def fake_repl(code):
    import ast
    module_ast = ast.parse(mycode)
    bytecode = compile(module_ast, '<fakefilename>', 'exec')
    global_ns = {}
    local_ns = {}
    exec(bytecode, global_ns, local_ns)
    return local_ns

fake_repl(mycode)
hey
Out[11]:
{'a': 1}

We don't show global_ns as it is huge, it will contain all that's availlable by default in Python. Let see where it fails if you use try a top-level async statement:

In [12]:
import ast
mycode = """
import aiohttp
await aiohttp.get('https://aip.github.com/')
"""

module_ast = ast.parse(mycode)
  File "<unknown>", line 3
    await aiohttp.get('https://aip.github.com/')
                ^
SyntaxError: invalid syntax

Ouch, so we can't even compile it. Let be smart can we get the inner code ? if we wrap in async-def ?

In [13]:
mycode = """
async def fake():
    import aiohttp
    await aiohttp.get('https://aip.github.com/')
"""
module_ast = ast.parse(mycode)
ast.dump(module_ast)
Out[13]:
"Module(body=[AsyncFunctionDef(name='fake', args=arguments(args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), body=[Import(names=[alias(name='aiohttp', asname=None)]), Expr(value=Await(value=Call(func=Attribute(value=Name(id='aiohttp', ctx=Load()), attr='get', ctx=Load()), args=[Str(s='https://aip.github.com/')], keywords=[])))], decorator_list=[], returns=None)])"
In [14]:
ast.dump(module_ast.body[0])
Out[14]:
"AsyncFunctionDef(name='fake', args=arguments(args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), body=[Import(names=[alias(name='aiohttp', asname=None)]), Expr(value=Await(value=Call(func=Attribute(value=Name(id='aiohttp', ctx=Load()), attr='get', ctx=Load()), args=[Str(s='https://aip.github.com/')], keywords=[])))], decorator_list=[], returns=None)"

As a reminder, as AST stands for Abstract Syntax Tree, you may construct an AST which is not a valid Python, program, like an if-else-else. AST tree can be modified. What we are interested in it the body of the function, which itself is the first object of a dummy module:

In [15]:
body = module_ast.body[0].body
body
Out[15]:
[<_ast.Import at 0x105d503c8>, <_ast.Expr at 0x105d50438>]

Let's pull out the body of the function and put it at the top level of a newly created module:

In [16]:
async_mod = ast.Module(body)
ast.dump(async_mod)
Out[16]:
"Module(body=[Import(names=[alias(name='aiohttp', asname=None)]), Expr(value=Await(value=Call(func=Attribute(value=Name(id='aiohttp', ctx=Load()), attr='get', ctx=Load()), args=[Str(s='https://aip.github.com/')], keywords=[])))])"

Mouahahahahahahahahah, you managed to get a valid top-level async ast ! Victory is yours !

In [17]:
bytecode = compile(async_mod, '<fakefile>', 'exec')
  File "<fakefile>", line 4
SyntaxError: 'await' outside function

Grumlgrumlgruml. You haven't said your last word. Your going to take your revenge later. Let's see waht we can do in Part II, not written yet.

Changing ByteStr REPR

A recent rebutal against Python 3 was recently written by the (in)famous Zed Shaw, with many responses to various arguments and counter arguments.

One particular topic which caught my eye was the bytearray vs unicodearray debate. I'll try explicitely avoid the term str/string/bytes/unicode naming as it is (IMHO) confusing, but that's a debate for another time. If one pay attention to above debates, you might see that there are about two camps:

  • bytearray and unicodearray are two different things, and we should never convert from one to the other. (that's rought the Pro-Python-3 camp)
  • bytearray and unicodearray are similar enough in most cases that we should do the magic for users.

I'm greatly exagerating here and the following is neither for one side or another, I have my personal preference of what I think is good, but that's irrelevant for now. Note that both sides argue that their preference is better for beginners.

You can often find posts trying to explain the misconception string/str/bytes, like this one which keep insisting on the fact that str in python 3 is far different from bytes.

The mistake in the REPR

I have one theory that the bytes/str issue is not in their behavior, but in their REPR. The REPR is in the end the main informatin communication channel between the object and the brain of the programmer, user. Also, Python "ducktyped", and you have to admit that bytes and str kinda look similar when printed, so assuming they should behave in similar way is not far fetched. I'm not saying that user will conciously assume bytes/str are the same. I'm saying that human brain inherently may do such association.

From the top of your head, what does requests.get(url).content returns ?

In [1]:
import requests_cache
import requests
requests_cache.install_cache('cachedb.tmp')
In [2]:
requests.get('http://swapi.co/api/people/1').content
Out[2]:
b'{"name":"Luke Skywalker","height":"172","mass":"77","hair_color":"blond","skin_color":"fair","eye_color":"blue","birth_year":"19BBY","gender":"male","homeworld":"http://swapi.co/api/planets/1/","films":["http://swapi.co/api/films/6/","http://swapi.co/api/films/3/","http://swapi.co/api/films/2/","http://swapi.co/api/films/1/","http://swapi.co/api/films/7/"],"species":["http://swapi.co/api/species/1/"],"vehicles":["http://swapi.co/api/vehicles/14/","http://swapi.co/api/vehicles/30/"],"starships":["http://swapi.co/api/starships/12/","http://swapi.co/api/starships/22/"],"created":"2014-12-09T13:50:51.644000Z","edited":"2014-12-20T21:17:56.891000Z","url":"http://swapi.co/api/people/1/"}'

... bytes...

I'm pretty sure you glanced ahead in this post and probaly thought it was "Text", even probably in this case Json. It might be invalid Json, I'm pretty sure you cannot tell.

Why does it returns bytes ? Because it could fetch an image:

In [3]:
requests.get('https://avatars0.githubusercontent.com/u/335567').content[:200]
Out[3]:
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x01\xcc\x00\x00\x01\xcc\x08\x06\x00\x00\x00X\xdb\x98\x86\x00\x00 \x00IDATx\xda\xac\xbdy\x93\x1b\xb9\xb2\xf6\xf7K\x00\xb5\x90\xbdH\xa3\x99\xb9s7\xbf\xf1:\x1c\x0e/\xdf\xff\xdb8\xec\xb0}\xd79g4Rw\xb3IV\x15\x80\xf4\x1f@\xedUl\xea\\w\x84\xa65-6Y\x85\x02ry\xf2\xc9'\xa5\xfe\x9f\xfeGE\x04#\x821\x061\x16c\x0c\xc6XD\x0c\x02\xa0\x8a\x8a\x801\xa4\x1f\x08\x880\xfdRUD\x04\xd5\xfe\xff#6z\x8c*\xaa\x82\x88\xe0C \x84@\xf7~\xa6yy\xc5=>Q>~\xe6\xe1\xf3g~\xfd\xa7\x7f\xc28\x07\xb6\x00\x84h-\x88A1(\xe0U\xd2\xfb\xb8t\r1("

And if you decode the first request ?

In [4]:
requests.get('http://swapi.co/api/people/2').content.decode()
Out[4]:
'{"name":"C-3PO","height":"167","mass":"75","hair_color":"n/a","skin_color":"gold","eye_color":"yellow","birth_year":"112BBY","gender":"n/a","homeworld":"http://swapi.co/api/planets/1/","films":["http://swapi.co/api/films/5/","http://swapi.co/api/films/4/","http://swapi.co/api/films/6/","http://swapi.co/api/films/3/","http://swapi.co/api/films/2/","http://swapi.co/api/films/1/"],"species":["http://swapi.co/api/species/2/"],"vehicles":[],"starships":[],"created":"2014-12-10T15:10:51.357000Z","edited":"2014-12-20T21:17:50.309000Z","url":"http://swapi.co/api/people/2/"}'

Well that looks the same (except leading b...). Go explain a beginner that the 2 above are totally different things, while they already struggle with 0 base indexing, iterators, and the syntax of the language.

Changing the repr

Lets revert the repr of bytesarray to better represent what they are. IPython allows to change object repr easily:

In [5]:
text_formatter = get_ipython().display_formatter.formatters['text/plain']
In [6]:
def _print_bytestr(arg, p, cycle):
    p.text('<BytesBytesBytesBytesBytes>')        
text_formatter.for_type(bytes, _print_bytestr)
Out[6]:
<function IPython.lib.pretty._repr_pprint>
In [7]:
requests.get('http://swapi.co/api/people/4').content
Out[7]:
<BytesBytesBytesBytesBytes>

Make a usefull repr

<bytesbytesbytes> may not an usefull repr, so let's try to make a repr, that:

  • Convey bytes are, in genral not text.
  • Let us peak into the content to guess what it is
  • Push the user to .decode() if necessary.

Generally in Python objects have a repr which start with <, then have the class name, a quoted representation of the object, and memory location of the object, a closing >.

As the _quoted representation of the object may be really long, we can ellide it.

A common representation of bytes could be binary, but it's not really compact. Hex, compact but more difficult to read, and make peaking at the content hart when it could be ASCII. So let's go with ASCII reprentation where we escape non ASCII caracterd.

In [8]:
ellide = lambda s: s if (len(s) < 75) else  s[0:50]+'...'+s[-16:]
In [9]:
def _print_bytestr(arg, p, cycle):
    p.text('<bytes '+ellide(repr(arg))+' at {}>'.format(hex(id(arg))))       
text_formatter.for_type(bytes, _print_bytestr)
Out[9]:
<function __main__._print_bytestr>
In [10]:
requests.get('http://swapi.co/api/people/12').content
Out[10]:
<bytes b'{"name":"Wilhuff Tarkin","height":"180","mass":"...pi/people/12/"}' at 0x107299228>
In [11]:
requests.get('http://swapi.co/api/people/12').content.decode()
Out[11]:
'{"name":"Wilhuff Tarkin","height":"180","mass":"unknown","hair_color":"auburn, grey","skin_color":"fair","eye_color":"blue","birth_year":"64BBY","gender":"male","homeworld":"http://swapi.co/api/planets/21/","films":["http://swapi.co/api/films/1/","http://swapi.co/api/films/6/"],"species":["http://swapi.co/api/species/1/"],"vehicles":[],"starships":[],"created":"2014-12-10T16:26:56.138000Z","edited":"2014-12-20T21:17:50.330000Z","url":"http://swapi.co/api/people/12/"}'

Advantage: It is not gobbledygook anymore when getting binary resources !

In [12]:
requests.get('https://avatars0.githubusercontent.com/u/335567').content
Out[12]:
<bytes b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x01\...0IEND\xaeB`\x82' at 0x107e0c000>

Remapping notebook shortcuts

As Jupyter notebook run in a browser for technical and practical reasons we only have a limited number of shortcuts available and choices need to be made. Often this choices may conflict with browser shortcut, and you might need to remap it.

Today I was inform by Stefan van Der Walt that Cmd-Shift-P conflict for Firefox. It is mapped both to open the Command palette for the notebook and open a new Private Browsing window.

Using Private Browsing windows is extremely useful. When developing a website you might want to look at it without being logged in, and with an empty cache. So let see how we can remap the Jupyter notebook shortcut.

TL; DR;

Use the following in your ~/.jupyter/custom/custom.js :

require(['base/js/namespace'], function(Jupyter){
  // we might want to but that in a callback on wait for 
  // en even telling us the ntebook is ready.
  console.log('== remaping command palette shortcut ==')
  // note that meta is the command key on mac.
  var source_sht = 'meta-shift-p'
  var target_sht = 'meta-/'
  var cmd_shortcuts = Jupyter.keyboard_manager.command_shortcuts;
  var action_name = cmd_shortcuts.get_shortcut(source_sht)
  cmd_shortcuts.add_shortcut(target_sht, action_name)
  cmd_shortcuts.remove_shortcut(source_sht)
  console.log('== ', action_name, 'remaped from', source_sht, 'to', target_sht )
})

details

We need to use require and register a callback once the notebook is loaded:

require(['base/js/namespace'], function(Jupyter){
  ...
})

Here we grab the main namespace and name it Jupyter.

Then get the object that hold the various shortcuts: var cmd_shortcuts = Jupyter.keyboard_manager.command_shortcuts.

Shortcuts are define by sequence on keys with modifiers. Modifiers are dash-separated (need to be pressed at the same time). Sequence are comma separated. Example quiting in vim would be esc,;,w,q, in emacs ctrl-x,ctrl-c.

Here we want to unbind meta-shift-p (p is lowercase despite shift being pressed) and bind meta-/ (The shortcut Stefan wants). Note that meta- is the command key on mac.

We need to get the current command bound to this shortcut (cmd_shortcuts.get_shortcut(source_sht)). You could hardcode the name of the command but it may change a bit depending on notebook version (this is not yet public API). Here it is jupyter-notebook:show-command-palette.

You now bind it to your new shortcut:

cmd_shortcuts.add_shortcut('meta-/', action_name)

And finally unbind the original one

cmd_shortcuts.remove_shortcut('meta-shift-p')

UI reflect your changes !

If you open the command palette, you should see that the Show command palette command now display Command-/ as its shortcut !

Future

We are working on an interface to edit shortcuts directly from within the UI and not to have to write a single line of code !

Questions, feedback and fixes welcomed

Viridisify

Viridisify

As usual this is available and as been written as a jupyter notebook if you like to play with the code feel free to fork it.


The jet colormap (AKA "rainbow") is ubiquitous, there are a lot of controverse as to wether it is (from far) the best one. And better options have been designed.

The question is, if you have a graph that use a specific colormap, and you would prefer for it to use another one; what do you do ?

Well is you have th eunderlying data that's easy, but it's not always the case.

So how to remap a plot which has a non perceptually uniform colormap using another ? What's happend if yhere are encoding artificats and my pixels colors are slightly off ?

I came up with a prototype a few month ago, and was asked recently by @stefanv to "correct" a animated plot of huricane Matthew, where the "jet" colormap seem to provide an illusion of growth:

https://twitter.com/stefanvdwalt/status/784429257556492288

Let's see how we can convert a "Jet" image to a viridis based one. We'll first need some assumptions:

  • This assume that you "know" the initial color map of a plot, and that the emcoding/compressing process of the plot will not change the colors "too much".
  • There are pixels in the image which are not part of the colormap (typically text, axex, cat pictures....)

We will try to remap all the pixels that fall not "too far" from the initial colormap to the new colormap.

In [1]:
%matplotlib inline
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
In [2]:
import matplotlib.colors as colors
In [3]:
!rm *.png *.gif out*
rm: output.gif: No such file or directory

I used the following to convert from mp4 to image sequence (8 fps determined manually). Sequence of images to video, and video to gif (quality is better than to gif dirrectly):

$ ffmpeg -i INPUT.mp4 -r 8 -f image2 img%02d.png
$ ffmpeg -framerate 8 -i vir-img%02d.png -c:v libx264 -r 8 -pix_fmt yuv420p out.mp4
$ ffmpeg -i out.mp4  output.gif
In [4]:
%%bash
ffmpeg -i input.mp4 -r 8 -f image2 img%02d.png -loglevel panic

Let's take our image without the alpha channel, so only the first 3 components:

In [5]:
import matplotlib.image as mpimg
img = mpimg.imread('img01.png')[:,:,:3]
In [6]:
fig, ax = plt.subplots()
ax.imshow(img)
fig.set_figheight(10)
fig.set_figwidth(10)

As you can see it does use "Jet" (most likely),

let's look at the repartitions of pixels on the RGB space...

In [7]:
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

import matplotlib.pyplot as plt
In [8]:
def rep(im, cin=None, sub=128):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    pp = im.reshape((-1,3)).T[:,::300]
    
    if cin:
        cmapin = plt.get_cmap(cin)
        cmap256 = colors.makeMappingArray(sub, cmapin)[:, :3].T
        ax.scatter(cmap256[0], cmap256[1], cmap256[2], marker='.', label='colormap', c=range(sub), cmap=cin, edgecolor=None)
    
    ax.scatter(pp[0], pp[1], pp[2], c=pp.T, marker='+')
    
    ax.set_xlabel('R')
    ax.set_ylabel('G')
    ax.set_zlabel('B')
    ax.set_title('Color of pixels')
    if cin:
        ax.legend()
    return ax
    
ax = rep(img)

We can see a specific clusers of pixel, let's plot the location of our "Jet" colormap and a diagonal of "gray". We can guess the effect of various compressions artifacts have jittered the pixels slightly away from their original location.

Let's look at where the jet colormap is supposed to fall:

In [9]:
rep(img, 'jet')
Out[9]:
<matplotlib.axes._subplots.Axes3DSubplot at 0x111c9cc88>

Ok, that's pretty accurate, we also see that our selected graph does nto use the full extent of jet.

in order to find all the pixels that uses "Jet" efficiently we will use scipy.spatial.KDTree in the colorspace. In particular we will subsample the initial colormap in sub=256 subsamples, and collect only pixels that are within d=0.2 of this subsample, and map each of these pixels to the closer subsample.

As we know the subsampling of the initial colormap, we can also determine the output colors.

The Pixels that are "too far" from the pixels of the colormap are keep unchanged.

increasing 256 to higher value will give a smoother final colormap.

In [10]:
from scipy.spatial import cKDTree
In [11]:
def convert(sub=256, d=0.2, cin='jet', cout='viridis', img=img, show=True):
    viridis = plt.get_cmap(cout)
    cmapin = plt.get_cmap(cin)
    cmap256 = colors.makeMappingArray(sub, cmapin)[:, :3]
    original_shape = img.shape
    img_data = img.reshape((-1,3))
    
    # this will efficiently find the pixels "close" to jet
    # and assign them to which point (from 1 to 256) they are on the colormap.
    K = cKDTree(cmap256)
    res = K.query(img_data, distance_upper_bound=d)
    
    indices = res[1]
    l = len(cmap256)
    indices = indices.reshape(original_shape[:2])
    remapped = indices

    indices.max()

    mask = (indices == l)

    remapped = remapped / (l-1)
    mask = np.stack( [mask]*3, axis=-1)

    # here we add only these pixel and plot them again with viridis.
    blend = np.where(mask, img, viridis(remapped)[:,:,:3])
    if show:
        fig, ax = plt.subplots()
        fig.set_figheight(10)
        fig.set_figwidth(10)
        ax.imshow(blend)
    return blend
In [12]:
res = convert(img=img)
rep(res)
Out[12]:
<matplotlib.axes._subplots.Axes3DSubplot at 0x113791278>