model.py

import os
import sys
import subprocess

# Get script's absolute directory path
script_dir = os.path.dirname(os.path.abspath(__file__))
requirements_path = os.path.join(script_dir, "requirements.txt")

# Install with path validation
try:
    subprocess.check_call([
        sys.executable,
        "-m", 
        "pip",
        "install",
        "-r",
        requirements_path
    ])
except subprocess.CalledProcessError as e:
    print(f"Installation failed with error code {e.returncode}")
    sys.exit(1)
except FileNotFoundError:
    print(f"requirements.txt not found at: {requirements_path}")
    sys.exit(1)


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