diff --git a/src/main/java/io/bioimage/modelrunner/model/processing/Processing.java b/src/main/java/io/bioimage/modelrunner/model/processing/Processing.java index fb3020f6..5c78f454 100644 --- a/src/main/java/io/bioimage/modelrunner/model/processing/Processing.java +++ b/src/main/java/io/bioimage/modelrunner/model/processing/Processing.java @@ -34,6 +34,7 @@ import io.bioimage.modelrunner.tensor.Tensor; import net.imglib2.type.NativeType; import net.imglib2.type.numeric.RealType; +import net.imglib2.util.Cast; /** * Class that executes the pre- or post-processing associated to a given tensor. @@ -132,15 +133,18 @@ List> preprocess(List> tensorList){ */ public & NativeType, R extends RealType & NativeType> List> preprocess(List> tensorList, boolean inplace) { + List> outputs = new ArrayList>(); + if (preMap.entrySet().size() == 0) return Cast.unchecked(tensorList); for (Entry> ee : this.preMap.entrySet()) { Tensor tt = tensorList.stream().filter(t -> t.getName().equals(ee.getKey())).findFirst().orElse(null); if (tt == null) continue; - ee.getValue().forEach(trans -> { - trans.run(tt, inplace); - }); + for (TransformationInstance trans : ee.getValue()) { + List> outList = trans.run(tt, inplace); + outputs.addAll(outList); + } } - return null; + return outputs; } /** @@ -174,14 +178,17 @@ List> postprocess(List> tensorList){ */ public & NativeType, R extends RealType & NativeType> List> postprocess(List> tensorList, boolean inplace) { + List> outputs = new ArrayList>(); + if (postMap.entrySet().size() == 0) return Cast.unchecked(tensorList); for (Entry> ee : this.postMap.entrySet()) { Tensor tt = tensorList.stream().filter(t -> t.getName().equals(ee.getKey())).findFirst().orElse(null); if (tt == null) continue; - ee.getValue().forEach(trans -> { - trans.run(tt, inplace); - }); + for (TransformationInstance trans : ee.getValue()) { + List> outList = trans.run(tt, inplace); + outputs.addAll(outList); + } } - return null; + return outputs; } }