13
13
import os
14
14
import shutil
15
15
import tempfile
16
+ from ast import Bytes
16
17
from pathlib import Path
17
18
from typing import Dict , Optional
18
19
29
30
trimesh , _ = optional_import ("trimesh" )
30
31
31
32
import monai .deploy .core as md
32
- from monai .deploy .core import DataPath , ExecutionContext , Image , InputContext , IOType , Operator , OutputContext
33
+ from monai .deploy .core import ExecutionContext , Image , InputContext , IOType , Operator , OutputContext
33
34
34
35
__all__ = ["STLConversionOperator" , "STLConverter" ]
35
36
36
37
37
38
@md .input ("image" , Image , IOType .IN_MEMORY )
38
- @md .output ("stl_output" , DataPath , IOType .DISK )
39
+ @md .output ("stl_output" , Bytes , IOType .IN_MEMORY ) # Only available when run as non-leaf operator
39
40
# nibabel is required by the dependent class STLConverter.
40
41
@md .env (
41
42
pip_packages = ["numpy>=1.21" , "nibabel >= 3.2.1" , "numpy-stl>=2.12.0" , "scikit-image>=0.17.2" , "trimesh>=3.8.11" ]
42
43
)
43
44
class STLConversionOperator (Operator ):
44
- """Converts volumetric image to surface mesh in STL format, file output only."""
45
+ """Converts volumetric image to surface mesh in STL format, file output only.
46
+
47
+ Only when used as a non-leaf operator is the output of STL binary stored in memory idenfied by the output label.
48
+ If a file path is provided, the STL binary will be saved in the the application's output folder of the current run.
49
+ """
45
50
46
51
def __init__ (
47
52
self , output_file = None , class_id = None , is_smooth = True , keep_largest_connected_component = True , * args , ** kwargs
@@ -59,16 +64,17 @@ def __init__(
59
64
self ._class_id = class_id
60
65
self ._is_smooth = is_smooth
61
66
self ._keep_largest_connected_component = keep_largest_connected_component
62
- self ._output_file = output_file if output_file and len (output_file ) > 0 else None
67
+ self ._output_file = output_file if output_file and len (str ( output_file ) ) > 0 else None
63
68
64
69
self ._converter = STLConverter (* args , ** kwargs )
65
70
66
71
def compute (self , op_input : InputContext , op_output : OutputContext , context : ExecutionContext ):
67
72
"""Gets the input (image), processes it and sets results in the output.
68
73
69
74
When used in a leaf operator, this function cannot set its output as in-memory object due to
70
- current limitation, and only file output, for DataPath IOType_DISK, will be saved in the
71
- op_output path, which is mapped to the application's output path by the execution engine.
75
+ current limitation.
76
+ If a file path is provided, the STL binary will be saved in the the application's output
77
+ folder of the current run.
72
78
73
79
Args:
74
80
op_input (InputContext): An input context for the operator.
@@ -80,20 +86,21 @@ def compute(self, op_input: InputContext, op_output: OutputContext, context: Exe
80
86
if not input_image :
81
87
raise ValueError ("Input is None." )
82
88
83
- op_output_config = op_output .get ()
84
- if self ._output_file and len (self ._output_file ) > 0 :
85
- # The file output folder is either the op_output or app's output depending on output types.
86
- output_folder = (
87
- op_output_config .path if isinstance (op_output_config , DataPath ) else context .output .get ().path
88
- )
89
- self ._output_file = output_folder / self ._output_file
90
- self ._output_file .parent .mkdir (exist_ok = True )
91
- self ._logger .info (f"Output will be saved in file { self ._output_file } ." )
89
+ # Use the app's current run output folder as parent to the STL output path.
90
+ if self ._output_file and len (str (self ._output_file )) > 0 :
91
+ _output_file = context .output .get ().path / self ._output_file
92
+ _output_file .parent .mkdir (parents = True , exist_ok = True )
93
+ self ._logger .info (f"Output will be saved in file { _output_file } ." )
92
94
93
- stl_bytes = self ._convert (input_image , self . _output_file )
95
+ stl_bytes = self ._convert (input_image , _output_file )
94
96
95
- if not isinstance (op_output_config , DataPath ):
96
- op_output .set (stl_bytes )
97
+ try :
98
+ # TODO: Need a way to find if the operator is run as leaf node in order to
99
+ # avoid setting in_memory object.
100
+ if self .op_info .get_storage_type ("output" , "stl_output" ) == IOType .IN_MEMORY :
101
+ op_output .set (stl_bytes )
102
+ except Exception as ex :
103
+ self ._logger .warn (f"In_memory output cannot be used when run as non-leaf operator. { ex } " )
97
104
98
105
def _convert (self , image : Image , output_file : Optional [Path ] = None ):
99
106
"""
@@ -152,12 +159,8 @@ def convert(
152
159
if not image or not isinstance (image , Image ):
153
160
raise ValueError ("image is not a Image object." )
154
161
155
- if not isinstance (output_file , Path ):
156
- raise ValueError ("output_file is not a Path" )
157
-
158
- # Ensure output file's folder exists
159
- if output_file .parent :
160
- output_file .parent .mkdir (exist_ok = True )
162
+ if isinstance (output_file , Path ):
163
+ output_file .parent .mkdir (parents = True , exist_ok = True )
161
164
162
165
s_image = self .SpatialImage (image )
163
166
nda = s_image .image_array
0 commit comments