Isoline Filter with ImageVisual

contour.py

import sys
from vispy import scene
from vispy import app
from vispy import visuals
from vispy.util.filter import gaussian_filter
from vispy.visuals.shaders import Function, FunctionChain
from vispy.color import Color
import numpy as np
from vispy.io import load_data_file, read_png

_null_color_transform = 'vec4 pass(vec4 color) { return color; }'
_c2l = 'float cmap(vec4 color) { return (color.r + color.g + color.b) / 3.; }'

_texture_bicubic = """
    // coefficients for the cubic polynomial
    vec4 c0 = vec4(-1.0,  3.0, -3.0,  1.0 ) /  6.0;
    vec4 c1 = vec4( 3.0, -6.0,  0.0,  4.0 ) /  6.0;
    vec4 c2 = vec4(-3.0,  3.0,  3.0,  1.0 ) /  6.0;
    vec4 c3 = vec4( 1.0,  0.0,  0.0,  0.0 ) /  6.0;

    vec4 cubic(vec4 var, vec4 p0, vec4 p1, vec4 p2, vec4 p3 ) {
    // return cubic polynomial
        return  p0 * dot( c0, var) +
            p1 * dot( c1, var) +
            p2 * dot( c2, var) +
            p3 * dot( c3, var);
    }

    // bicubic interpolation
    //

    vec4 texture_lookup(vec2 texcoord) {
        if(texcoord.x < 0.0 || texcoord.x > 1.0 ||
        texcoord.y < 0.0 || texcoord.y > 1.0) {
            discard;
        }

        vec2 uv = texcoord;
        vec2 shape = $shape;
        vec2 p = shape * uv - 0.5;
        vec2 a = fract(p);
        vec2 xy = floor(p);

        float x = a.x;
        float x2 = x * x;
        float x3 = x * x2;
        vec4 varx = vec4(x3, x2, x, 1.0);

        vec4 r0 = cubic( varx,
                        texture2D($texture, (xy + vec2(-1, -1))/shape),
                        texture2D($texture, (xy + vec2(0, -1))/shape),
                        texture2D($texture, (xy + vec2(1, -1))/shape),
                        texture2D($texture, (xy + vec2(2, -1))/shape));

        vec4 r1 = cubic( varx,
                        texture2D($texture, (xy + vec2(-1, 0))/shape),
                        texture2D($texture, (xy + vec2(0, 0))/shape),
                        texture2D($texture, (xy + vec2(1, 0))/shape),
                        texture2D($texture, (xy + vec2(2, 0))/shape));

        vec4 r2 = cubic( varx,
                        texture2D($texture, (xy + vec2(-1, 1))/shape),
                        texture2D($texture, (xy + vec2(0, 1))/shape),
                        texture2D($texture, (xy + vec2(1, 1))/shape),
                        texture2D($texture, (xy + vec2(2, 1))/shape));

        vec4 r3 = cubic( varx,
                        texture2D($texture, (xy + vec2(-1, 2))/shape),
                        texture2D($texture, (xy + vec2(0, 2))/shape),
                        texture2D($texture, (xy + vec2(1, 2))/shape),
                        texture2D($texture, (xy + vec2(2, 2))/shape));

        x = a.y;
        x2 = x * x;
        x3 = x * x2;
        vec4 vary = vec4(x3, x2, x, 1.0);

        return cubic( vary, r0, r1, r2, r3);
    }"""


class Isoline(object):
    def __init__(self, level=10., width=1.0, color='black'):
        self.shader = Function("""
            void isoline() {
                // function taken from glumpy/examples/isocurves.py
                // and extended to have level, width and color as parameters
                
                // Extract data value
                float value = gl_FragColor.r;

                // setup lw, aa
                float antialias = 1.0;
                float linewidth = $isowidth + antialias;

                // "middle" contour(s) dividing upper and lower half
                // but only if isolevel is even
                if( mod($isolevel,2.0) == 0.0 ) {
                    if( length(value - 0.5) < 0.5 / $isolevel)
                        linewidth = linewidth * 2;
                }

                // Trace contour
                float v  = $isolevel * value - 0.5;
                float dv = linewidth/2.0 * fwidth(v);
                float f = abs(fract(v) - 0.5);
                float d = smoothstep(-dv,+dv,f);
                float t = linewidth/2.0 - antialias;
                d = abs(d)*linewidth/2.0 - t;

                if( d < - linewidth ) {
                    d = 1.0;
                } else  {
                     d /= antialias;
                }

                vec4 bg = $color_transform1(gl_FragColor);
                //bg = gl_FragColor;

                vec4 fc = vec4($isocolor.rgb, 0);

                if (d < 1.) {
                    //fc *= d;
                    fc.a = 1-d;
                }

                gl_FragColor = mix(bg, fc, fc.a);
            }
        """)
        self.level = level
        self.width = width
        self.color = color

    @property
    def level(self):
        return self._level

    @level.setter
    def level(self, l):
        self._level = l
        self.shader['isolevel'] = l

    @property
    def width(self):
        return self._width

    @width.setter
    def width(self, w):
        self._width = w
        self.shader['isowidth'] = w

    @property
    def color(self):
        return self._color

    @color.setter
    def color(self, c):
        self._color = c
        self.shader['isocolor'] = Color(c).rgba

    def _attach(self, visual):
        visual._get_hook('frag', 'post').add(self.shader())
        fun = FunctionChain(None, [Function(_c2l),
                                       Function(visual._cmap.glsl_map)])
        self.shader['color_transform1'] = fun
        visual.shared_program.frag['color_transform'] = Function(_null_color_transform)

canvas = scene.SceneCanvas(keys='interactive')
canvas.size = 800, 600
canvas.show()

# Set up a viewbox to display the image with interactive pan/zoom
view = canvas.central_widget.add_view()

# Create the image
interpolation = 'bicubic'
np.random.seed(1000)
img_data = np.random.normal(size=(100, 100), loc=50, scale=150)
img_data = gaussian_filter(img_data, (4, 4, 0)).astype(np.float32)

#img_data = np.zeros(25).reshape((5, 5)).astype(np.float32)
#img_data[1:4, 1::2] = 0.5
#img_data[1::2, 2] = 0.5
#img_data[2, 2] = 1.0

#img_data = read_png(load_data_file('mona_lisa/mona_lisa_sm.png'))
#image = scene.visuals.Image(img_data, interpolation=interpolation,
#                            parent=view.scene, method='subdivide')

img_data1 = np.zeros((100, 100))
img_data1[0,0] = -0.5
img_data1[99,99] = +0.5
image = scene.visuals.Image(img_data, interpolation=interpolation,
                            parent=view.scene, method='subdivide')

image._interpolation_fun['bicubic2D'] = Function(_texture_bicubic)
image._interpolation_names = image._interpolation_names + ['bicubic2D']
image.interpolation_filters = image._interpolation_names

#image = scene.visuals.Contour(img_data, interpolation='bicubic1',
#                            parent=view.scene, method='subdivide',
#                            levels=2, width=1.0, color='yellow')

image.transform = visuals.transforms.STTransform(translate=(0, 0, 0.5))

levels = np.linspace(img_data.min(), img_data.max(), num=13, endpoint=True)[1:-1]
color_lev = 'white'
# Create isocurve, make a child of the image to ensure the two are always
# aligned.
#curve = scene.visuals.Isocurve(img_data, levels=levels, color_lev=color_lev,
#                              parent=view.scene, method='subdivide')

iso = Isoline(level=10, width=3., color='black')

canvas.title = 'Spatial Filtering using %s Filter' % interpolation

# Set 2D camera (the camera will scale to the contents in the scene)
view.camera = scene.PanZoomCamera(aspect=1)
# flip y-axis to have correct aligment
view.camera.flip = (0, 1, 0)
view.camera.set_range()

# get interpolation functions from Image
names = image.interpolation_functions
names = list(names)
names.sort()
print(names)
act = 17

act = -1

level = 2
first = True


# Implement key presses
@canvas.events.key_press.connect
def on_key_press(event):
    global act, level, first, interpolation
    if event.key in ['Left', 'Right']:
        if event.key == 'Right':
            step = 1
        else:
            step = -1
        act = (act + step) % len(names)
        interpolation = names[act]
        image.interpolation = interpolation

    if event.key in ['Up', 'Down']:
        if first:
            image.attach(iso)
            first = False
        if event.key == 'Up':
            level += 1
        else:
            if level > 1:
                level -= 1
        levels = np.linspace(img_data.min(), img_data.max(), num=level+1, endpoint=True)[1:-1]
        #curve.levels = levels
        iso.level = level
    canvas.title = 'Spatial Filtering using %s Filter - Isoline %d level' % (interpolation, level)
    canvas.update()


if __name__ == '__main__' and sys.flags.interactive == 0:
    app.run()