How to send images using FastAPI and WebSocket + Three.js

WebSockets is an advanced technology that makes it possible to open a two-way interactive communication session between a browser and a server. Then, you can send messages to a server and receive event-driven responses without having to poll the server for a reply 🤓

Now, is this good? In some cases there is no way to monitor if WebSockets based services are working. This might be bad for production environments. Also, there are vulnerabilities problems (well, it's over the Internet). In any case, let me set my problem:

Send images from a 3D volume to a browser's client in real-time.
The volume is loaded in a webserver 😜

A Websocket establishes a bi-directional communication, meaning both the client and server can send/receive data. The common example is a chat, but in this case I don't need messages! I need images.

It is important to note that WebSockets is a TCP-based protocol, like http but using ws or wss (further information in the RFC 6455)

What I am expecting?

Given a volume of dimensions w x h x d, the browser's client could render an image of size w x h at position i, where 0 <= i < d. The ith position is set using a slider.

**This is what I need to develop**: render images over Three.js, and images could change using a slider

via GIPHY

Web Server

To begin with, I created a web server using FastAPI (Python), which supports WebSocket. The documentation is exceptional and besides all the code required for a webserver in FastAPI, this is the code using WebSockets:

from fastapi import WebSocket

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    print("started")
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            print(f"received: {int(data)}")
            index = int(data)
            image = get_image(volume, index)
            await websocket.send_bytes(image)
    except Exception as e:
        print(e)
    finally:
        websocket.close()

There are some things to note in the code above:

Asynchronous function (async)
The keyword await to pass control back
Text is received but bytes are sent. The received text represents the value from a slider put in the client.
The variable volume is a 3D NumPy array, and the function get_image returns the image at the position index (in Base64).

The function to convert the NumPy array into a Base64, these are functions:

def image_to_base64(img: np.ndarray) -> bytes:
    """ Given a numpy 2D array, returns a JPEG image in base64 format """

    # using opencv 2, there are others ways
    img_buffer = cv2.imencode('.jpg', img)[1]
    return base64.b64encode(img_buffer).decode('utf-8')
    
def get_image(volume, index: int):
    image = volume[:, :, index]
    return image_to_base64(image)

For now, the code for the server is done. Just started using for instance uvicorn. My server is in the port 8080. Let's go to the client side.

Client

For the client-side, it is necessary an HTML page and the JavaScript code. To do not enter in details about the HTML code; it is important the canvas tag. At this point, it is possible to use several things related to any Javascript library or framework. To simplify things, we need to add something like this:

<canvas id="myCanvas" width="800" height="600"></canvas>

Notice that the name of canvas is important, which will be used in the Javascript code. The canvas is required to create a context for Three.js. Three.js is a cross-browser JavaScript library an API used to create and display animated 3D computer graphics in a web browser using WebGL.

To ease the display, instead to use a classic img tag, I preferred to employ a 3D plane into a canvas with the ability to zoom in/out, rotate, translate and more. Three.js, now transferred to NPM as three package might be used incorporate in your code as a module or via CDN (content delivery network). In my case, I used the CDN from Skypack. The code associated with the Three.js variables are defined as follow:

import * as THREE from 'https://cdn.skypack.dev/three@0.129.0/build/three.module.js';
import { OrbitControls } from 'https://cdn.skypack.dev/three@0.129.0/examples/jsm/controls/OrbitControls.js';

const canvas = document.querySelector('#myCanvas');
const renderer = new THREE.WebGLRenderer({canvas});
const scene = new THREE.Scene();
const camera = new THREE.PerspectiveCamera( 75, canvas.width / canvas.height, 0.1, 1000 );

Notice the querySelector command which is equivalent to the getElementById to locate a component, in our case: myCanvas.

For the WebSocket, we need to add the following:

const URL = "ws://192.168.1.173:8080/ws";
var socket = new WebSocket(URL);
socket.binaryType = "arraybuffer";

To make simpler the transmission, we need to set the property arraybuffer to control the data received over the WebSocket. Thus, we can explore the available functions in the MDN Web Docs.

To receive and obtain data, the definition of required functions might be as follow:

socket.onopen = function(e) {
    console.log("[open] Connection established");
};

socket.onmessage = function(event) {
    console.log("[message] Data received from server:");

    const arrayBuffer = event.data;

    image_Slice.src = 'data:image/jpg;base64,' + arrayBuffer;
    console.log("size= "+ arrayBuffer.length);
};

socket.onclose = function(event) {
    if (event.wasClean) {
        console.log("[close] Connection closed cleanly, code=${event.code} reason=${event.reason}");
    } else {
        console.log("[close] Connection died");
    }
};

socket.onerror = function(error) {
    console.log("[error] ${error.message}");
};

The function to receive the data using WebSocket is the onmessage. Notice that image_Slice represents a variable of type Image created as
let image_Slice = new Image();
Next, that image is placed over a Three.js texture using a single MeshLambertMaterial. Remember that received data are bytes in Base64 format.

On the other side, to data sent is a string that represents the value from the slider. In that case, the send function is placed in the onchange function of the element slider:

slider.onchange =()=>{
    console.log("slider:" + slider.value)
    socket.send(slider.value)
    var output = document.getElementById("myRangeValue");
    output.innerHTML = "index: " +slider.value;
}

For this point, we are done 😎. To recap: a volume is loaded in a web server-side and send to a client using the Base64 encoding. This image will be displayed into a 3D plane in Three.js. To test the advantage of WebSockets, a slider was added to send which slice from the volume client needs. Here is a video to see how it looks.