streamtrace 0.1.0

Streamtrace SDK — decorator-based ML pipeline contracts for clinical AI

streamtrace

Decorator-based ML pipeline contracts for clinical AI. Define your pipeline once with typed decorators, get a validated execution DAG and a deployable frontend app.

Installation

pip install streamtrace

Install extras for the output types your pipeline produces:

pip install streamtrace[medical]    # NIfTI + DICOM (nibabel, pydicom, numpy)
pip install streamtrace[imaging]    # images + matplotlib (Pillow, matplotlib, numpy)
pip install streamtrace[plotting]   # interactive plots (matplotlib, plotly, bokeh)
pip install streamtrace[data]       # tabular output (pandas, numpy)
pip install streamtrace[mesh]       # 3D geometry (trimesh, numpy)
pip install streamtrace[all]        # everything

Quick start

import streamtrace as st

@st.app(title="Cardiac Segmentation", version="1.0.0",
        docker_base="pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime",
        requirements=["nibabel", "scipy"])
class CardiacSegmentation:

    @st.input(title="CT Scan",
              widget=st.FileInput(accepted_types=[st.InputDataType.NIFTI]),
              returns="scan")
    def load_scan(self, path):
        import nibabel as nib
        return nib.load(path).get_fdata()

    @st.input(title="Threshold",
              widget=st.TextInput(default="0.5", placeholder="0.0 – 1.0"),
              returns="threshold")
    def set_threshold(self, value):
        return float(value)

    @st.preprocess(returns="normalized")
    def normalize(self, scan):
        lo, hi = scan.min(), scan.max()
        return (scan - lo) / (hi - lo + 1e-8)

    @st.infer(returns="raw_mask",
              device="cuda",
              weights_path="checkpoints/best.pth")
    def segment(self, normalized, threshold):
        import torch
        model = torch.load("checkpoints/best.pth", map_location="cpu")
        model.eval()
        with torch.no_grad():
            t = torch.FloatTensor(normalized).unsqueeze(0).unsqueeze(0)
            pred = model(t).squeeze().numpy()
        return (pred > threshold).astype("uint8")

    @st.postprocess(returns="mask")
    def clean(self, raw_mask):
        from scipy.ndimage import binary_fill_holes
        return binary_fill_holes(raw_mask).astype("uint8")

    @st.output(title="Segmentation",
               widget=st.FileOutput(data_type=st.OutputDataType.NIFTI,
                                    filename="segmentation.nii.gz"))
    def save(self, mask):
        import nibabel as nib
        return nib.Nifti1Image(mask, affine=None)

Pipeline contract

A Streamtrace pipeline is a class decorated with @st.app. Each method in the class maps to one node in the execution DAG via its decorator. Wiring is automatic: a method's parameter names are matched against the returns= alias of other nodes.

@st.input  →  @st.preprocess  →  @st.infer  →  @st.postprocess  →  @st.output

Decorators

Decorator	Purpose	Key args
@st.app	Marks the pipeline class	title, version, description, docker_base, requirements
@st.input	User-provided value (file, text, etc.)	title, widget, required, returns
@st.preprocess	Data preparation before inference	title, returns
@st.infer	Model inference step	title, returns, device, weights_path
@st.postprocess	Clean-up after inference	title, returns
@st.output	Final (or intermediate) result	title, widget, returns, intermediate

Wiring via returns

@st.input(returns="scan")       # produces key "scan"
def load_scan(self, path): ...

@st.preprocess(returns="normalized")
def normalize(self, scan): ...  # "scan" resolves to load_scan's output

@st.infer(returns="mask")
def segment(self, normalized): ...  # "normalized" resolves to normalize's output

If returns is omitted the method name is used as the output key.

Input widgets

Widget	Description
FileInput(accepted_types=[...], max_size_mb=500)	File upload
TextInput(default="", placeholder="", max_length=None)	Single-line text

Accepted file types — InputDataType.ANY, .NIFTI, .DICOM, .PNG, .NUMPY

Output widgets

Widget	Description
FileOutput(data_type=OutputDataType.File, filename=None)	Generic file download
ImageOutput(caption="", max_size_mb=50)	PNG / matplotlib figure
PlotOutput(interactive=True)	Plotly / Bokeh / matplotlib
MetricOutput(label="", precision=3)	Numeric value or dict of metrics

Output data types — File, Image, Plot, NIFTI, DICOM, Mesh, CSV, JSON, Text, Metric

Intermediate outputs are uploaded to the frontend while the pipeline is still running:

@st.output(title="Slice Preview", widget=st.ImageOutput(), intermediate=True)
def preview(self, normalized): ...

DAG inspection

from streamtrace import build_dag

dag = build_dag(CardiacSegmentation)

# Validate wiring before deploying
errors = dag.validate()
if errors:
    for e in errors:
        print(e)

# Topological execution order
for node in dag.topological_sort():
    print(node.phase, "→", node.name, "produces", node.output_key)

# JSON schema (used by streamtrace push)
import json
print(json.dumps(dag.to_schema(), indent=2))

Streamtrace CLI

The SDK is the contract layer. The streamtrace CLI handles the rest:

streamtrace init       # surveys your codebase and classifies existing code
streamtrace configure  # builds and validates a complete executable pipeline
streamtrace push       # deploys to Streamtrace infrastructure

License

MIT