TkTorchWindow.py

# Copyright 2020 Erik Härkönen. All rights reserved.
# This file is licensed to you under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You may obtain a copy
# of the License at http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software distributed under
# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
# OF ANY KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.

import tkinter as tk
import numpy as np
import time
from contextlib import contextmanager
import pycuda.driver
from pycuda.gl import graphics_map_flags
from glumpy import gloo, gl
from pyopengltk import OpenGLFrame
import torch
from torch.autograd import Variable

# TkInter widget that can draw torch tensors directly from GPU memory

@contextmanager
def cuda_activate(img):
    """Context manager simplifying use of pycuda.gl.RegisteredImage"""
    mapping = img.map()
    yield mapping.array(0,0)
    mapping.unmap()

def create_shared_texture(w, h, c=4,
        map_flags=graphics_map_flags.WRITE_DISCARD,
        dtype=np.uint8):
    """Create and return a Texture2D with gloo and pycuda views."""
    tex = np.zeros((h,w,c), dtype).view(gloo.Texture2D)
    tex.activate() # force gloo to create on GPU
    tex.deactivate()
    cuda_buffer = pycuda.gl.RegisteredImage(
        int(tex.handle), tex.target, map_flags)
    return tex, cuda_buffer

# Shape batch as square if possible
def get_grid_dims(B):
    S = int(B**0.5 + 0.5)
    while B % S != 0:
        S -= 1
    return (B // S, S)

def create_gl_texture(tensor_shape):
    if len(tensor_shape) != 4:
        raise RuntimeError('Please provide a tensor of shape NCHW')
    
    N, C, H, W = tensor_shape

    cols, rows = get_grid_dims(N)
    tex, cuda_buffer = create_shared_texture(W*cols, H*rows, 4)

    return tex, cuda_buffer

# Create window with OpenGL context
class TorchImageView(OpenGLFrame):
    def __init__(self, root = None, show_fps=True, **kwargs):
        self.root = root or tk.Tk()
        self.width = kwargs.get('width', 512)
        self.height = kwargs.get('height', 512)
        self.show_fps = show_fps
        self.pycuda_initialized = False
        self.animate = 0 # disable internal main loop
        OpenGLFrame.__init__(self, root, **kwargs)

    # Called by pyopengltk.BaseOpenGLFrame
    # when the frame goes onto the screen
    def initgl(self):
        if not self.pycuda_initialized:
            self.setup_gl(self.width, self.height)
            self.pycuda_initialized = True
        
        """Initalize gl states when the frame is created"""
        gl.glViewport(0, 0, self.width, self.height)
        gl.glClearColor(0.0, 0.0, 0.0, 0.0)
        self.dt_history = [1000/60]
        self.t0 = time.time()
        self.t_last = self.t0
        self.nframes = 0

    def setup_gl(self, width, height):
        # setup pycuda and torch
        import pycuda.gl.autoinit
        import pycuda.gl

        assert torch.cuda.is_available(), "PyTorch: CUDA is not available"
        print('Using GPU {}'.format(torch.cuda.current_device()))
        
        # Create tensor to be shared between GL and CUDA
        # Always overwritten so no sharing is necessary
        dummy = torch.cuda.FloatTensor((1))
        dummy.uniform_()
        dummy = Variable(dummy)
        
        # Create a buffer with pycuda and gloo views, using tensor created above
        self.tex, self.cuda_buffer = create_gl_texture((1, 3, width, height))
        
        # create a shader to program to draw to the screen
        vertex = """
        uniform float scale;
        attribute vec2 position;
        attribute vec2 texcoord;
        varying vec2 v_texcoord;
        void main()
        {
            v_texcoord = texcoord;
            gl_Position = vec4(scale*position, 0.0, 1.0);
        } """
        fragment = """
        uniform sampler2D tex;
        varying vec2 v_texcoord;
        void main()
        {
            gl_FragColor = texture2D(tex, v_texcoord);
        } """
        # Build the program and corresponding buffers (with 4 vertices)
        self.screen = gloo.Program(vertex, fragment, count=4)
        
        # NDC coordinates:         Texcoords:          Vertex order,
        # (-1, +1)       (+1, +1)   (0,0)     (1,0)    triangle strip:
        #        +-------+               +----+          1----3
        #        |  NDC  |               |    |          |  / | 
        #        | SPACE |               |    |          | /  |
        #        +-------+               +----+          2----4
        # (-1, -1)       (+1, -1)   (0,1)     (1,1)
        
        # Upload data to GPU
        self.screen['position'] = [(-1,+1), (-1,-1), (+1,+1), (+1,-1)]
        self.screen['texcoord'] = [(0,0), (0,1), (1,0), (1,1)]
        self.screen['scale'] = 1.0
        self.screen['tex'] = self.tex

    # Don't call directly, use update() instead
    def redraw(self):
        t_now = time.time()
        dt = t_now - self.t_last
        self.t_last = t_now

        self.dt_history = ([dt] + self.dt_history)[:50]
        dt_mean = sum(self.dt_history) / len(self.dt_history)

        if self.show_fps and self.nframes % 60 == 0:
            self.master.title('FPS: {:.0f}'.format(1 / dt_mean))

    def draw(self, img):
        assert len(img.shape) == 4, "Please provide an NCHW image tensor"
        assert img.device.type == "cuda", "Please provide a CUDA tensor"

        if img.dtype.is_floating_point:
            img = (255*img).byte()
        
        # Tile images
        N, C, H, W = img.shape

        if N > 1:
            cols, rows = get_grid_dims(N)
            img = img.reshape(cols, rows, C, H, W)
            img = img.permute(2, 1, 3, 0, 4) # [C, rows, H, cols, W]
            img = img.reshape(1, C, rows*H, cols*W)

        tensor = img.squeeze().permute(1, 2, 0).data # CHW => HWC
        if C == 3:
            tensor = torch.cat((tensor, tensor[:,:,0:1]),2) # add the alpha channel
            tensor[:,:,3] = 1 # set alpha
        
        tensor = tensor.contiguous()

        tex_h, tex_w, _ = self.tex.shape
        tensor_h, tensor_w, _ = tensor.shape

        if (tex_h, tex_w) != (tensor_h, tensor_w):
            print(f'Resizing texture to {tensor_w}*{tensor_h}')
            self.tex, self.cuda_buffer = create_gl_texture((N, C, H, W)) # original shape
            self.screen['tex'] = self.tex

        # copy from torch into buffer
        assert self.tex.nbytes == tensor.numel()*tensor.element_size(), "Tensor and texture shape mismatch!"
        with cuda_activate(self.cuda_buffer) as ary:
            cpy = pycuda.driver.Memcpy2D()
            cpy.set_src_device(tensor.data_ptr())
            cpy.set_dst_array(ary)
            cpy.width_in_bytes = cpy.src_pitch = cpy.dst_pitch = self.tex.nbytes//tensor_h
            cpy.height = tensor_h
            cpy(aligned=False)
            torch.cuda.synchronize()
        
        # draw to screen
        self.screen.draw(gl.GL_TRIANGLE_STRIP)

    def update(self):
        self.update_idletasks()
        self.tkMakeCurrent()
        self.redraw()
        self.tkSwapBuffers()

# USAGE:
# root = tk.Tk()
# iv = TorchImageView(root, width=512, height=512)
# iv.pack(fill='both', expand=True)
# while True:
#     iv.draw(nchw_tensor)
#     root.update()
#     iv.update()