Skip to content

Commit df8de2d

Browse files
committed
In the end, there can be only one file_name
1 parent 7d5cba8 commit df8de2d

File tree

1 file changed

+1
-7
lines changed

1 file changed

+1
-7
lines changed

geoengine/ml.py

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ class MlModelConfig:
1919
'''Configuration for an ml model'''
2020
name: str
2121
metadata: MlModelMetadata
22-
file_name: str = "model.onnx"
2322
display_name: str = "My Ml Model"
2423
description: str = "My Ml Model Description"
2524

@@ -41,7 +40,7 @@ def register_ml_model(onnx_model: ModelProto,
4140

4241
with geoengine_openapi_client.ApiClient(session.configuration) as api_client:
4342
with tempfile.TemporaryDirectory() as temp_dir:
44-
file_name = Path(temp_dir) / model_config.file_name
43+
file_name = Path(temp_dir) / model_config.metadata.file_name
4544

4645
with open(file_name, 'wb') as file:
4746
file.write(onnx_model.SerializeToString())
@@ -74,11 +73,6 @@ def check_data_type(data_type: TypeProto, expected_type: RasterDataType, prefix:
7473
raise InputException(f'Model {prefix} type `{elem_type_str}` does not match the '
7574
f'expected type `{expected_type}`')
7675

77-
for domain in onnx_model.opset_import:
78-
if domain.domain != '':
79-
continue
80-
if domain.version != 9:
81-
raise InputException('Only ONNX models with opset version 9 are supported')
8276

8377
model_inputs = onnx_model.graph.input
8478
model_outputs = onnx_model.graph.output

0 commit comments

Comments
 (0)