Skip to content

Cleanup Tasks, PredictionWriter, and to_onnx.py

Samuel Van Stroud requested to merge svanstro/salt:svanstro/onnx into main

Description

This MR cleans up the prediction writer and onnx export code by moving functionality to the relevant Task classes.

Main changes:

  • Tasks are now decorated with the model name
  • Tasks should implement their own output names as properties
  • Tasks should implement get_h5() and get_onnx() functions, which are used in the prediction writer and onnx export
  • Cleanup PredictionWriter and to_onnx.py using this

Other changes:

  • Update atlas-ftag-tools to fix precision of prediction writer outputs

Work towards #64

cc @npond and @backes, I'd appreciate a second look here and @backes should probably check I haven't broken anything. @npond in the longer term we should see if we can fit the maskformer outputs into this scheme to clean things up further.

Review checklist:

  • CI Passing
  • Comments addressed
  • Source branch is up to date with target
Edited by Samuel Van Stroud

Merge request reports