Running additional thread(s) for Tensorflow predictions within Slicer

Development
, , ,
1

Hello,

I am using Slicer to perform real-time (and also offline - but that is less affected by this issue) AI-based segmentation of tracked ultrasound images, while concurrently recording the images using Sequences; and in some cases - reconstructing/rendering the segmentations into the 3D viewer simultaneously.

The issue that I have is that when doing the real-time segmentation and reconstruction, the recording frame-rate within Sequences drops as (what I believe to be happening, anyway is) everything runs in the main thread, meaning that a new frame won’t be recorded in Sequences until the Tensorflow AI model has done its prediction on that frame, and the resulting segmentation is added to the volume reconstruction.

What I would like to do is have the recording running on the main thread, and have the AI predictions (and possibly the volume reconstruction) run in a separate thread or process to prevent them from restricting one another.

Any suggestions on best practices (CLI module, using Python’s multi-threading or some other PyPI package for threading, etc.), or any direction on how this could be done would be much appreciated.

Thank you!

2 Likes
2

This might be a use case for the SlicerProcesses approach. Basically it spawns PythonSlicer as worker processes and communicates with them through pickled dictionaries (e.g. of numpy arrays and metadata). PythonSlicer has the same environment as Slicer, so you can import the same packages but you don’t have access to the Slicer qt app or vtk rendering context of the application.

You might also look at how SimpleFilters releases the python GIL to use python threading.

As you work on this it would be great if you could share a sample script or module of whatever approach works for you as a reference for others who want to solve the same problem in the future.

2 Likes
3

Hi @pieper - Those both seem like simple enough solutions!

Will give them a try and come back with questions or some sample code to share here for others.

4

To minimize interference between the GUI application and processing, you can run the prediction in a separate process (even on a separate computer) that communicates with the data acquisition server (Plus) and the GUI (Slicer) using OpenIGTLink. The processing server can send and receive OpenIGTLink messages using pyigtl in a simple loop. The processing server can receive the input images directly from Plus and only send stream the results to Slicer.

 Plus ->------------->--------------->- Slicer
         \                         /           
          \                       /            
           -> Processing server ->
1 Like
5

For anyone looking to see roughly how I’ve ended up implementing this, see the gist here, which demonstrates a live and offline prediction using the SlicerProcesses module @pieper proposed. I’ve inserted some comments to clarify some choices, but this will need to be tailored to a specific application/module when used.

Our previous application which did live predictions, volume reconstruction and recording with Sequences on streamed US images from PLUS obtained ~10FPS recording speed while reconstructing those 10FPS in real-time (on fairly modern hardware). Using the above methods, we obtain ~25-30FPS.

For the offline predictions, which previously took 10-15 seconds and froze the application; things still take the same amount of time, but there is no “locking up” of Slicer while the predictions occur in the background.

I’ve listed some of the key changes/notes below:

The offline prediction is very similar to the examples provided in that repository, but

  • You must pass a path to the model (or just its weights) to the process as TF models are not picklable
  • I am using fairly large volumes in the offline predictions, I write to a temporary file instead of purely using stdin and stdout

The live prediction is a bit more tricky;

  • Instead of writing just once, we need to be able to write continuously to stdin’s buffer
  • Then, update the data attribute of the Process object, and call the method which sends the new data
  • Additionally, with that data - provide an “on” or “off” byte and the size of the data, so we know that we’ve received everything properly in the script
  • Instead of receiving everything at once when the process finishes, connect a callback to the Process which happens each time stdout is written to
  • Terminate by switching the active byte “off”

Hopefully this helps others looking to do similar things in the future.

2 Likes
6

Nice work :+1: Thanks for posting the follow up info.