import numpy as np
import triton_python_backend_utils as pb_utils
from omnicloudmask import predict_from_array

class TritonPythonModel:
    def initialize(self, args):
        pass

    def execute(self, requests):
        responses = []
        for request in requests:
            # Get input tensor
            input_tensor = pb_utils.get_input_tensor_by_name(request, "input_array")
            input_array = input_tensor.as_numpy()

            # Perform inference
            pred_mask = predict_from_array(input_array)

            # Create output tensor
            output_tensor = pb_utils.Tensor(
                "output_mask", 
                pred_mask.astype(np.uint8)
            )
            responses.append(pb_utils.InferenceResponse([output_tensor]))
        return responses

    def finalize(self):
        pass