inference.py

"""
The following is a simple example algorithm.

It is meant to run within a container.

To run the container locally, you can call the following bash script:

  ./do_test_run.sh

This will start the inference and reads from ./test/input and writes to ./test/output

To save the container and prep it for upload to Grand-Challenge.org you can call:

  ./do_save.sh

Any container that shows the same behaviour will do, this is purely an example of how one COULD do it.

Reference the documentation to get details on the runtime environment on the platform:
https://grand-challenge.org/documentation/runtime-environment/

Happy programming!
"""

from pathlib import Path
import json
import tifffile 
import numpy as np
import torch


# Constants for the location of the input and output files, please do not modify! 
INPUT_PATH = Path("/input/images/image-stack-unstructured-noise")
OUTPUT_PATH = Path("/output/images/image-stack-denoised")
OUTPUT_PATH.mkdir(parents=True, exist_ok=True)


# ============================================================================
# USER-CUSTOMIZABLE FUNCTIONS
# Modify these functions to implement your own inference pipeline
# ============================================================================

def load_model():
    """
    Load your model. You have two options for the model path:
    1. Save model as a part of the Docker-container image in the resources/ directory.
        It will be available at the /opt/app/resources directory at runtime.
    2. Upload them as a separate tarball to Grand Challenge (go to your Algorithm > Models). 
        The resources in the tarball will be extracted to /opt/ml/model directory at runtime.
    """
    # Option 1: part of the Docker-container image: resources/
    model_path = Path("/opt/app/resources/my_model.pth")

    # Option 2: upload them as a separate tarball to Grand Challenge (go to your Algorithm > Models). 
    # The resources in the tarball will be extracted to `model_dir` at runtime.
    # Example: model_path = Path("/opt/ml/model/my_model.pth")
    
    print(f"Loading model: {model_path}")
    return torch.jit.load(model_path)


def read_image(image_path: Path) -> np.ndarray:
    """
    Read and preprocess input image.
    Modify this function to implement your own image loading and preprocessing pipeline.
    """
    print(f"Reading image: {image_path}")
    input_array = tifffile.imread(image_path)
    input_array = input_array.astype(np.float32)
    print(f"Loaded image shape: {input_array.shape}")
    
    return input_array


def run_inference(model, input_tensor):
    """
    Run inference on the input tensor.
    Modify this function to implement your own inference logic.
    """
    print("Running inference...")
    print(f"Input shape: {input_tensor.shape}")
    input_tensor = torch.from_numpy(input_tensor)
    with torch.no_grad():
        output = model(input_tensor).squeeze().cpu().numpy()
    print(f"Output shape: {output.shape}")
    return output


def save_output(array, output_path):
    """
    Save the processed array.
    """
    print(f"Saving output to: {output_path}")
    with tifffile.TiffWriter(output_path) as out:
        out.write(
            array,
            resolutionunit=2 # This flag is important for the GC to process the output correctly! 
        )


def inference_handler():
    """
    Main handler for processing images with unstructured noise.
    This is where you should implement your main inference pipeline.
    """
    # Show torch cuda info
    _show_torch_cuda_info()

    # Load your model
    model = load_model()

    # Load and prepare input
    input_files = sorted(INPUT_PATH.glob("*.tif"))
    print(f"Reading input files: {input_files}")

    for input_file in input_files:
        input_array = read_image(input_file)

        # Run inference
        result = run_inference(model, input_array)

        # Save output
        output_path = OUTPUT_PATH / input_file.name
        save_output(result, output_path)

    return 0


# ============================================================================
# UTILITY FUNCTIONS
# These functions handle the interface with the Grand Challenge platform
# ============================================================================

def run():
    # The key is a tuple of the slugs of the input sockets
    interface_key = get_interface_key()

    print(f"Interface key: {interface_key}")

    handler = inference_handler

    # Call the handler
    return handler()


def get_interface_key():
    # The inputs.json is a system generated file that contains information about
    # the inputs that interface with the algorithm
    inputs = load_json_file(INPUT_PATH.parent.parent / "inputs.json")
    socket_slugs = [sv["interface"]["slug"] for sv in inputs]
    return tuple(sorted(socket_slugs))


def load_json_file(location):
    # Reads a json file
    with open(location, "r") as f:
        return json.loads(f.read())


def _show_torch_cuda_info():
    print("=+=" * 10)
    print("Collecting Torch CUDA information")
    print(f"Torch CUDA is available: {(available := torch.cuda.is_available())}")
    if available:
        print(f"\tnumber of devices: {torch.cuda.device_count()}")
        print(f"\tcurrent device: { (current_device := torch.cuda.current_device())}")
        print(f"\tproperties: {torch.cuda.get_device_properties(current_device)}")
    print("=+=" * 10)


if __name__ == "__main__":
    raise SystemExit(run())