truthdotphd commited on
Commit
36d88d2
·
verified ·
1 Parent(s): 19b46c5

Create model.py

Browse files
Files changed (1) hide show
  1. model.py +28 -0
model.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import triton_python_backend_utils as pb_utils
3
+ from omnicloudmask import predict_from_array
4
+
5
+ class TritonPythonModel:
6
+ def initialize(self, args):
7
+ pass
8
+
9
+ def execute(self, requests):
10
+ responses = []
11
+ for request in requests:
12
+ # Get input tensor
13
+ input_tensor = pb_utils.get_input_tensor_by_name(request, "input_array")
14
+ input_array = input_tensor.as_numpy()
15
+
16
+ # Perform inference
17
+ pred_mask = predict_from_array(input_array)
18
+
19
+ # Create output tensor
20
+ output_tensor = pb_utils.Tensor(
21
+ "output_mask",
22
+ pred_mask.astype(np.uint8)
23
+ )
24
+ responses.append(pb_utils.InferenceResponse([output_tensor]))
25
+ return responses
26
+
27
+ def finalize(self):
28
+ pass