Tree computational structure with different inputs and outputs

Hi everyone,

I’m new to rust (and heavily typed languages in general) coming from python. I’m interested in making an image processing tool for a side project, but I’m having some difficulty of coming up with the design of the structure in rust.

The outcome I want is to be able to define a tree of processes to perform on images, like applying histograms transformations, cropping, merging two images, averaging a collection of images, etc. So a user can make a workflow on their images. I guess the ultimate goal would be to have a UI where people can add nodes of processes and connect them, but I might be getting ahead of myself.

From the examples I can find online this would be okay if all the processes worked on the same data, say something like

trait Process {
    fn process(&self, input: Image) -> Image;
}

Then I would be able to create a tree of nodes using a Vec and can keep track of the inputs and outputs. Maybe something like this:

struct Node {
    process: Box<dyn Process>,
    input: usize,
    output: usize,
}

struct Tree {
    nodes: Vec<Node>,
}

But my problem is that the I have processes which take different types of inputs, and numbers of inputs which also need to match up with the corresponding outputs of other nodes.
The types of inputs and outputs I want are greyscale images, rgb images, and collections of greyscale or rgb images.

Here’s some examples of processes I have:

  • FileImage which takes no inputs and loads a file from disk to an image
  • Histogram and cropping will take either a greyscale or rgb image and output a similar image
  • ChannelCombine takes exactly three greyscale inputs and maps them to an rgb image
  • some splitting which might take an image and output two images
  • AverageCollection will take a collection of greyscale or rgb images (potentially 1000’s of them saved to disk and average them together) and output a single image

I’ve been trying to implement something myself around this and keeping track of the types of inputs and outputs and having lots of match statements and raising errors, but I feel I’m just fighting the type system constantly.

I appreciate this is a long question and very specific but I’m just getting myself quite lost with it now and how someone who knows the language would start to tackle this problem, so any pointers would be appreciated! I can also provide better code examples of what I’ve tried before.

A few points to consider:

  1. You probably want a Directed Acyclic Graph and not a tree to represent your workflow.
  2. You can either roll your own graph data structure, or use an existing library like petgraph.
  3. Nodes in the graph represent sets of image files (inputs/outputs), and edges between nodes are image transformations.
  4. To validate your workflow, you can iterate through all edges and check that each edge's source and target nodes are compatible with the edge type. If not, you can create an informative error message.
  5. To execute the workflow, perform a topological sort on the graph (petgraph can do this) and process the nodes in sorted order.

For example, using petgraph, prototype code might look something like:

use petgraph::visit::EdgeRef;
use petgraph::{Directed, Graph};
use std::path::PathBuf;

#[derive(Debug)]
enum Process {
    FileImage,
    Histogram,
    ChannelCombine,
    SplitImage,
    AvgCollection,
}

#[derive(Debug)]
struct Image {
    file_path: PathBuf,
}

#[derive(Debug, Default)]
struct FileSet {
    images: Vec<Image>,
}

fn main() {
    // nodes are labelled with FileSets
    // edges are labelled with Processes
    let mut graph = Graph::<FileSet, Process, Directed>::new();

    // define input/output filesets
    let fileset1 = graph.add_node(FileSet {
        images: vec![Image {
            file_path: ["path", "to", "image1"].iter().collect(),
        }],
    });
    let fileset2 = graph.add_node(FileSet {
        images: vec![Image {
            file_path: ["path", "to", "image2"].iter().collect(),
        }],
    });
    let fileset3 = graph.add_node(FileSet {
        images: vec![Image {
            file_path: ["path", "to", "image3"].iter().collect(),
        }],
    });
    let fileset4 = graph.add_node(FileSet {
        images: vec![Image {
            file_path: ["path", "to", "image4"].iter().collect(),
        }],
    });
    let fileset5 = graph.add_node(FileSet {
        images: vec![Image {
            file_path: ["path", "to", "image5"].iter().collect(),
        }],
    });

    // add processes between filesets
    graph.add_edge(fileset1, fileset2, Process::AvgCollection);
    graph.add_edge(fileset1, fileset3, Process::Histogram);
    graph.add_edge(fileset3, fileset4, Process::ChannelCombine);
    graph.add_edge(fileset4, fileset5, Process::SplitImage);
    graph.add_edge(fileset3, fileset5, Process::FileImage);

    // validate edges
    for edge in graph.edge_references() {
        println!("This is my edge: {:#?}", edge);

        let input_files = &graph[edge.source()];
        let type_of_process = edge.weight();
        let output_files = &graph[edge.target()];

        // check the source and target nodes of this edge have
        // the right "shape"
        match type_of_process {
            &Process::FileImage => {}
            &Process::Histogram => {}
            &Process::ChannelCombine => {}
            &Process::SplitImage => {}
            &Process::AvgCollection => {}
        }

        // we can report an error if anything is wrong,
        // else proceed
    }
    // if there were no errors, perform a topological sort
    // (petgraph can do this)
}
2 Likes

Thank you very much for the detailed response, I hadn't thought of having the nodes in the graph as the images and then processes as the edges so that's really interesting! I think that will make the type system a lot clearer for me, and good to see how someone much more familiar with the language would approach it.

I'll have to play around with it myself for a while and make a mock up of all the different scenarios, as I'm still not quite clear on how something like the ChannelCombine will work where I have something like:

let red_image = graph.add_node(FileSet {...});
let green_image = graph.add_node(FileSet {...});
let blue_image = graph.add_node(FileSet {...});
let combined_image = graph.add_node(FileSet {...});

graph.add_edge(red_image, combined_image, Process:ChannelCombine);
graph.add_edge(green_image, combined_image, Process:ChannelCombine);
graph.add_edge(blue_image, combined_image, Process:ChannelCombine);

Because I have three separate images (which may have come from other processes) and want to combine them, but I can't do that with a single edge.

In that case, the nodes and edges should be swapped so that nodes are processes and edges are input data. You can then iterate over the nodes, using something like node_indices(), then make sure the incoming edges to each node -- obtained with edges_directed() -- are compatible with the node type.

(Note: you will need some kind of dummy Process node that has no input edges and just represents an existing image file on disk).

You may also need to add some kind of String name or numeric index to each FileSet. This will tell each target node which parameter the edge represents. This will be necessary if, for example, you need to subtract Image 1's pixel values from Image 2's, and you have to define which image is 1 and which is 2.

As a general piece of advice/hard-earned experience, it's usually helpful to start with a static picture of your key data structures/fields. It's much easier to reason about what your code is doing when you have a map of your data that isn't constantly moving around. When I read "keeping track of the types of inputs and outputs and having lots of match statements and raising errors", I guessed you'd skipped this stage and dived straight into code. :grinning:

I just wanted to follow up, that helped clear things up and how graph type structures are implemented in rust. Lots of learning still to do but thanks for the advice! Hopefully I have my image processing app in the not too distant future