Model Demo: UNET Tumor Segmentation

NOTE: The Jupyter Notebook below is included in the Chimera SDK and can be run interactively by running the following CLI command:

$ quadric sdk notebook

From the Jupyter Notebook window in your browser, select the notebook named /quadric/sdk-cli/examples/models/unet/unet.ipynb.

PATIENT_LIMIT = 1
LIMIT_PER_PATIENT = 3
TUMOR_THRESHOLD_FOR_ISS = 200
THREADS = 6
LIMITED_DATA = True
if LIMITED_DATA:
    !unzip -qo kaggle_cache.zip
else:
    !pip3 install kaggle -q
    #!KAGGLE_USERNAME=xxx KAGGLE_KEY=yyy kaggle datasets download -d mateuszbuda/lgg-mri-segmentation
    !kaggle datasets download -d mateuszbuda/lgg-mri-segmentation
    !if [ ! -d "kaggle_3m" ]; then unzip lgg-mri-segmentation.zip; fi

import torch

model = torch.hub.load(
    "mateuszbuda/brain-segmentation-pytorch",
    "unet",
    in_channels=3,
    out_channels=1,
    init_features=32,
    pretrained=True,
)

prefix = model.__class__.__name__
fp_onnx_name = f"{prefix}_float32.onnx"

Using cache found in /github/home/.cache/torch/hub/mateuszbuda_brain-segmentation-pytorch_master

## Init input tensor used for exporting in ONNX in NCHW

## assumptions for all imagenet-trained models. for user-supplied models change the following:
dataset_mean = [0.485, 0.456, 0.406]
dataset_std = [0.229, 0.224, 0.225]
dataset_input_size = (224, 224)  # an (W, H) tuple

## include quadric's cli helpers and instantiate a module to help
from sdk_cli.utils import model_helpers
import tvm.contrib.epu.chimera_job.constants as sdk_constants

mh = model_helpers.ModelHelper(dataset_input_size, dataset_mean, dataset_std)

x = torch.randn(1, 3, mh.size[1], mh.size[0], requires_grad=True)
## Export the model
torch.onnx.export(
    model,  # model being run
    x,  # model input (or a tuple for multiple inputs)
    fp_onnx_name,  # where to save the model (can be a file or file-like object)
    export_params=True,  # store the trained parameter weights inside the model file
    do_constant_folding=True,  # whether to execute constant folding for optimization
    opset_version=sdk_constants.DEFAULT_ONNX_OPSET,
    input_names=["input"],  # the model's input names
    output_names=["output"],
)  # the model's output names

print(f"Model exported to ONNX with name '{fp_onnx_name}'!")

Model exported to ONNX with name 'UNet_float32.onnx'!


/tmp/ipykernel_192659/2902265974.py:12: DeprecationWarning: Call to deprecated class ModelHelper. ('ModelHelper' class is being deprecated. Quadric APIs have been updated to use PyTorch datasets and transforms instead.) -- Deprecated since version 24.01.
  mh = model_helpers.ModelHelper(dataset_input_size, dataset_mean, dataset_std)

from torch.utils.data import DataLoader
from tqdm import tqdm
import torch
from dataset import BrainSegmentationDataset as Dataset
import onnxruntime as ort
from dataclasses import dataclass
import numpy as np
from skimage.io import imsave
from inference_onnx import (
    makedirs,
    data_loader,
    postprocess_per_volume,
    dsc_distribution,
    plot_dsc,
)
from utils import dsc, gray2rgb, outline
import os
from onnxruntime.quantization import CalibrationDataReader


@dataclass
class Arguments:
    onnx: str = fp_onnx_name
    images: str = "./kaggle_3m/"
    image_size: int = 224
    figure: str = "./dsc.png"
    predictions: str = "./predictions_onnx"
    batch_size: int = 1


args = Arguments()

makedirs(args)

loader = data_loader(args)

reading validation images...
preprocessing validation volumes...
cropping validation volumes...
padding validation volumes...
resizing validation volumes...
normalizing validation volumes...
done creating validation dataset

ort_sess = ort.InferenceSession(args.onnx)

input_list = []
pred_list = []
true_list = []

ort_data_reader = []
for i, data in tqdm(enumerate(loader)):
    x, y_true = data

    x_np = x.detach().cpu().numpy()
    ort_inp = {"input": x_np}
    ort_data_reader.append(ort_inp)

    y_pred_np = ort_sess.run(None, ort_inp)
    y_pred_np = np.array(y_pred_np[0])

    pred_list.extend([y_pred_np[s] for s in range(y_pred_np.shape[0])])

    y_true_np = y_true.detach().cpu().numpy()
    true_list.extend([y_true_np[s] for s in range(y_true_np.shape[0])])

    x_np = x.detach().cpu().numpy()
    input_list.extend([x_np[s] for s in range(x_np.shape[0])])

311it [00:10, 29.80it/s]

volumes = postprocess_per_volume(
    input_list,
    pred_list,
    true_list,
    loader.dataset.patient_slice_index,
    loader.dataset.patients,
)

for p in volumes:
    x = volumes[p][0]
    y_pred = volumes[p][1]
    y_true = volumes[p][2]
    for s in range(x.shape[0]):
        image = gray2rgb(x[s, 1])  # channel 1 is for FLAIR
        image = outline(image, y_pred[s, 0], color=[255, 0, 0])
        image = outline(image, y_true[s, 0], color=[0, 255, 0])
        filename = "{}-{}.png".format(p, str(s).zfill(2))
        filepath = os.path.join(args.predictions, filename)
        imsave(filepath, image)

/tmp/ipykernel_192659/1936848120.py:19: UserWarning: ./predictions_onnx/TCGA_DU_5854_19951104-33.png is a low contrast image
  imsave(filepath, image)
/tmp/ipykernel_192659/1936848120.py:19: UserWarning: ./predictions_onnx/TCGA_DU_5872_19950223-66.png is a low contrast image
  imsave(filepath, image)
/tmp/ipykernel_192659/1936848120.py:19: UserWarning: ./predictions_onnx/TCGA_DU_6399_19830416-50.png is a low contrast image
  imsave(filepath, image)
/tmp/ipykernel_192659/1936848120.py:19: UserWarning: ./predictions_onnx/TCGA_DU_5851_19950428-34.png is a low contrast image
  imsave(filepath, image)

from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import glob

img_width, img_height = 400, 400


def resize_img_to_array(img, img_shape=(244, 244)):
    img_array = np.array(img.resize(img_shape, Image.LANCZOS))

    return img_array


def image_grid(fn_images: list, text: list = [], top: int = 12, per_row: int = 4):
    """
    fn_images is a list of image paths.
    text is a list of annotations.
    top is how many images you want to display
    per_row is the number of images to show per row.
    """
    for i in range(len(fn_images[:top])):
        if i % 4 == 0:
            _, ax = plt.subplots(1, per_row, sharex="col", sharey="row", figsize=(20, 6))
        j = i % 4
        image = Image.open(fn_images[i])
        image = resize_img_to_array(image, img_shape=(img_width, img_height))
        ax[j].imshow(image)
        ax[j].axis("off")
        if text:
            ax[j].annotate(
                text[i],
                (0, 0),
                (0, -32),
                xycoords="axes fraction",
                textcoords="offset points",
                va="top",
            )


image_grid(
    fn_images=glob.glob(f"{args.predictions}/*.png"),
    text=glob.glob(f"{args.predictions}/*.png"),
)

import tvm.contrib.epu.chimera_job.quantize as quant
import onnx

inf_onnx = onnx.load(args.onnx)


class UnetDr(CalibrationDataReader):
    def __init__(self, array_of_inputs):
        self._enum_data_dicts = array_of_inputs
        self.reset()

    def get_next(self):
        return next(self.enum_data_dicts, None)

    def reset(self):
        self.enum_data_dicts = iter(self._enum_data_dicts)


unet_dr = UnetDr(ort_data_reader)

##optimize the model ahead of quantization

optimized_model_path, _ = quant.optimize_onnx(args.onnx)

quant_result = quant.quantize_float_onnx(optimized_model_path, unet_dr)

2026-06-19 04:18 - INFO - epu - quantize - Optimized model to opset
2026-06-19 04:18 - INFO - epu - quantize - Saved optimized model to UNet_float32_float32_opt.onnx
2026-06-19 04:18 - INFO - epu - quantize - Input shapes: [1, 3, 224, 224]. Input names: input
2026-06-19 04:18 - INFO - epu - quantize - Output shapes: [[1, 1, 224, 224]]. Output names: ['output']
2026-06-19 04:18 - DEBUG - epu - quantize - Full exclusion set for quantization: ['Softmax', 'Sigmoid', 'QuadricCustomOp']
2026-06-19 04:18 - DEBUG - epu - quantize - excl_nodes ['/Sigmoid']
2026-06-19 04:18 - INFO - epu - quantize - Quantization started...
WARNING:root:Please use QuantFormat.QDQ for activation type QInt8 and weight type QInt8. Or it will lead to bad performance on x64.
2026-06-19 04:18 - INFO - epu - quantize - Quantization done succesfully!
2026-06-19 04:18 - INFO - epu - quantize - ONNX full precision model size: 29.61 MB
2026-06-19 04:18 - INFO - epu - quantize - ONNX quantized model size: 7.45 MB
2026-06-19 04:18 - INFO - epu - quantize - Saved quantized model to /quadric/sdk-cli/examples/models/unet/UNet_opt_sym_int8_q.onnx
2026-06-19 04:18 - INFO - epu - quantize - Saved shape inferenced model to /quadric/sdk-cli/examples/models/unet/UNet_opt_sym_int8_q.onnx
2026-06-19 04:18 - INFO - epu - quantize - Checking for remaining FLOAT/FLOAT16 types.
2026-06-19 04:18 - INFO - epu - quantize - Model still has FLOAT/FLOAT16 types. Creating ranges for floating point tensors using calibration data


Custom quantization code for ConvTranspose
Custom quantization code for ConvTranspose
Custom quantization code for ConvTranspose
Custom quantization code for ConvTranspose


2026-06-19 04:19 - INFO - epu - quantize - Saved tensor ranges to /quadric/sdk-cli/examples/models/unet/UNet_opt_sym_int8_q.onnx.tranges

quant_result

QuantizeResult(qmodel_path='/quadric/sdk-cli/examples/models/unet/UNet_opt_sym_int8_q.onnx', tranges_path='/quadric/sdk-cli/examples/models/unet/UNet_opt_sym_int8_q.onnx.tranges')

from tvm.contrib.epu.chimera_job.chimera_job import ChimeraJob
from tvm.relay.backend.contrib.epu.util import logger as tvm_logger
import logging, warnings, sys

cgc_job = ChimeraJob(model_p=quant_result.qmodel_path, quiet_iss=False)
tvm_logger.setLevel(logging.INFO)
warnings.filterwarnings("ignore")
sys.tracebacklimit = 0
cgc_job.compile(quiet=True)

experiments = []
for patient, val in volumes.items():
    print(patient)
    relpath = f"./iss_results/{patient}"
    patient_iss_runs = 0
    for i in range(len(val[0])):
        patient_root = f"{patient}_{i}"
        inp = val[0][i]
        inp = np.expand_dims(inp, axis=0)
        onnx_output = ort_sess.run([], {"input": inp})[0][0]
        onnx_output = onnx_output.squeeze()
        if np.count_nonzero(np.round(onnx_output)) > TUMOR_THRESHOLD_FOR_ISS:
            experiment = {"patient": patient, "input": inp}
            experiments.append(experiment)
experiments = experiments[:THREADS]

TCGA_DU_5854_19951104
TCGA_CS_4944_20010208
TCGA_CS_4941_19960909
TCGA_DU_5872_19950223
TCGA_CS_6186_20000601
TCGA_CS_5397_20010315
TCGA_DU_5855_19951217
TCGA_CS_5393_19990606
TCGA_DU_6399_19830416
TCGA_DU_5851_19950428

ort_sess = ort.InferenceSession(args.onnx)
ort_sess_q = ort.InferenceSession(quant_result.qmodel_path)
from sdk_cli.visualizers.unet_brain import UnetBrainVisualizer

%matplotlib inline
from matplotlib.gridspec import GridSpec
import matplotlib.patches as mpatches
from sdk_cli.utils.model_helpers import calcualate_pixelwise_iou


## change these to run more patients and more images per patient
vis = UnetBrainVisualizer()


def run_iss_patient(experiment, unet_pixelvalue_threshold=0.75):
    print(f"patient: {experiment['patient']}")
    input = experiment["input"]
    onnx_output_q = ort_sess_q.run([], {"input": input})[0][0]
    onnx_output_q = onnx_output_q.squeeze()
    experiment["ort"] = onnx_output_q
    cgc_output = cgc_job.run_inference_harness(inputs={"input": input})["output"].reshape(
        onnx_output_q.shape
    )
    experiment["iss"] = cgc_output
    experiment["cgc_job"] = cgc_job
    experiment["iou"] = calcualate_pixelwise_iou(
        onnx_output_q > unet_pixelvalue_threshold,
        cgc_output > unet_pixelvalue_threshold,
    )
    return experiment


def plot_result(experiment):
    fig = plt.figure(constrained_layout=True, figsize=(12, 8))
    fig.suptitle(f"Patient {experiment['patient']}", y=0.92)
    gs = GridSpec(1, 2, figure=fig)
    ax = []
    ax.append(fig.add_subplot(gs[0, 0]))
    ax.append(fig.add_subplot(gs[0, 1]))

    ax[0].set_axis_off()
    ax[0].set_title("Brain Scan Data mapped to RGB")

    ax[1].set_title(f"Quantized ort and Chimera ISS. IoU: {experiment['iou']:0.2f}")
    legend = [
        mpatches.Patch(color="blue", label="ort"),
        mpatches.Patch(color="red", label="chimera iss"),
    ]
    ax[1].legend(handles=legend)
    input = experiment["input"]
    ax[0].imshow(UnetBrainVisualizer.input_img_to_display(input))
    onnx_output_q = experiment["ort"]
    cgc_output = experiment["iss"]
    image = vis.gray2rgb(input.squeeze()[1])
    image = vis.outline(image, onnx_output_q, color=[0, 0, 205])
    image = vis.outline(image, cgc_output, color=[255, 0, 0])
    experiment["image"] = image

    ax[1].imshow(image)

## output = run_iss_patient(experiments[0])

import multiprocess

with multiprocess.Pool(processes=THREADS) as pool:
    # process data in parallel
    results = pool.map(run_iss_patient, experiments)

patient: TCGA_CS_4944_20010208patient: TCGA_CS_4941_19960909

patient: TCGA_CS_4941_19960909
patient: TCGA_CS_4941_19960909
patient: TCGA_CS_4941_19960909
patient: TCGA_DU_5872_19950223


2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf8298205b0>
2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf829820460>
2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf829823100>
2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf8200919c0>
2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf820092410>
2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf829821360>
FILM 16/16: 100%|███████████████████████████████████████████████████| 16/16 [00:26<00:00,  1.67s/it]
2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf829820460>
2026-06-19 04:20 - INFO - epu - iss_testing - Started Executing Onnxruntime...
FILM 16/16: 100%|███████████████████████████████████████████████████| 16/16 [00:26<00:00,  1.67s/it]
2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf8298205b0>
2026-06-19 04:20 - INFO - epu - iss_testing - Started Executing Onnxruntime...
2026-06-19 04:20 - INFO - epu - iss_testing - Done 0:00:00.309790
FILM 16/16: 100%|███████████████████████████████████████████████████| 16/16 [00:27<00:00,  1.69s/it]
2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf820091810>
2026-06-19 04:20 - INFO - epu - iss_testing - Started Executing Onnxruntime...
FILM 16/16: 100%|███████████████████████████████████████████████████| 16/16 [00:27<00:00,  2.33s/it]2026-06-19 04:20 - INFO - epu - iss_testing - Done 0:00:00.322889
FILM 16/16: 100%|███████████████████████████████████████████████████| 16/16 [00:27<00:00,  1.69s/it]
2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf829823100>
2026-06-19 04:20 - INFO - epu - iss_testing - Started Executing Onnxruntime...
FILM 16/16: 100%|███████████████████████████████████████████████████| 16/16 [00:27<00:00,  1.71s/it]
2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf820093ee0>
2026-06-19 04:20 - INFO - epu - iss_testing - Started Executing Onnxruntime...
FILM 16/16: 100%|███████████████████████████████████████████████████| 16/16 [00:27<00:00,  1.71s/it]
2026-06-19 04:20 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7bf829821360>
2026-06-19 04:20 - INFO - epu - iss_testing - Started Executing Onnxruntime...
2026-06-19 04:20 - INFO - epu - iss_testing - Done 0:00:00.542599
2026-06-19 04:20 - INFO - epu - iss_testing - Done 0:00:00.526837
2026-06-19 04:20 - INFO - epu - iss_testing - Done 0:00:00.532566
2026-06-19 04:20 - INFO - epu - iss_testing - Done 0:00:00.435733

## print(output['input'])
for res in results:
    plot_result(res)

plt.show()

print(results[0]["cgc_job"])
results[0]["cgc_job"].plot_run_statistics()

╒═════════════════════╤════════════════════════════════════════════════════════════════════╕
│ Module Name         │ UNet_opt_sym_int8_q_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1 │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ ONNX File           │ /quadric/sdk-cli/examples/models/unet/UNet_opt_sym_int8_q.onnx     │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ Product Target      │ QC-U                                                               │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ Number of Cores     │ 1                                                                  │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ ISS Clock Frequency │ 1.700                                                              │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ L2M Size            │ 16MB                                                               │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ LRM Size            │ 4kB                                                                │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ External Read BW    │ 128GBps                                                            │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ External Write BW   │ 128GBps                                                            │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ MACS per PE         │ 16                                                                 │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ Max L2M             │ 10.608MB                                                           │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ Max LRM             │ 0.750kB                                                            │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ Max Temp Ext Bytes  │ 0.000MB                                                            │
├─────────────────────┼────────────────────────────────────────────────────────────────────┤
│ Network GMACs       │ 9.242                                                              │
╘═════════════════════╧════════════════════════════════════════════════════════════════════╛

╒════╤════════╤════════╤══════════════════╤══════════════════════════╤═══════╕
│    │ Type   │ Name   │ shape            │ type                     │ mse   │
╞════╪════════╪════════╪══════════════════╪══════════════════════════╪═══════╡
│  0 │ Input  │ input  │ [1, 3, 224, 224] │ tensor[FixedPoint32<27>] │ n/a   │
├────┼────────┼────────┼──────────────────┼──────────────────────────┼───────┤
│  1 │ Output │ output │ [1, 1, 224, 224] │ tensor[FixedPoint32<31>] │ 0.000 │
╘════╧════════╧════════╧══════════════════╧══════════════════════════╧═══════╛

Post-ISS Report 1.7 GHz ***
Fully placed-and-routed gate simulation: 
╒══════════════════════════════════╤═════════╕
│ Latency (ms)                     │ 0.92    │
├──────────────────────────────────┼─────────┤
│ FPS                              │ 1088.79 │
├──────────────────────────────────┼─────────┤
│ Average Power @ 3nm SSGNP (mW)   │ 2054.33 │
├──────────────────────────────────┼─────────┤
│ FPS per Watt @ 3nm SSGNP (FPS/W) │ 530.00  │
├──────────────────────────────────┼─────────┤
│ Ext Rd Bytes (MB)                │ 8.02    │
├──────────────────────────────────┼─────────┤
│ Ext Wr Bytes (MB)                │ 0.19    │
├──────────────────────────────────┼─────────┤
│ Avg Ext Rd BW (GBps)             │ 8.53    │
├──────────────────────────────────┼─────────┤
│ Avg Ext Wr BW (GBps)             │ 0.20    │
├──────────────────────────────────┼─────────┤
│ MAC Utilization                  │ 36.13%  │
╘══════════════════════════════════╧═════════╛
*** Data generated using 7nm SSGNP gatesim and scaled to 3nm

[SDK-CLI] : TotalCycles: 1,561,372
[SDK-CLI] : Executions/second: 1,088.79

compute      : ▇▇▇▇▇▇▇▇▇▇▇▇▇▇ 225.6K
data_array   : ▇▇▇▇▇▇▇▇▇ 151.006K
mac          : ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇ 791.045K
data_external: ▏ 7.985K
data_ocm     : ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇ 384.969K

for more information check run directory: /quadric/sdk-cli/examples/models/unet/ccl_build/UNet_opt_sym_int8_q_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1/run/20260619_042022_3dcaa4


2026-06-19 04:20 - INFO - epu - chimera_job - Combined plots generated and saved to: 
/quadric/sdk-cli/examples/models/unet/ccl_build/UNet_opt_sym_int8_q_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1/run/20260619_042022_3dcaa4/data/UNet_opt_sym_int8_q_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1.combined.png





'/quadric/sdk-cli/examples/models/unet/ccl_build/UNet_opt_sym_int8_q_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1/run/20260619_042022_3dcaa4/data'