#!/usr/bin/env python

__doc__ = """
Simple example script showing how to use the Sybil library locally to predict risk scores for a set of DICOM files.
"""

import sybil
from sybil import visualize_attentions

from utils import get_demo_data


def main():
    # Load a trained model
    model = sybil.Sybil("sybil_ensemble")

    dicom_files = get_demo_data()

    # Get risk scores
    serie = sybil.Serie(dicom_files)
    print(f"Processing {len(dicom_files)} DICOM files")
    prediction = model.predict([serie], return_attentions=True)
    scores = prediction.scores

    print(f"Risk scores: {scores}")

    # Visualize attention maps
    output_dir = "sybil_attention_output"

    print(f"Writing attention images to {output_dir}")
    series_with_attention = visualize_attentions(
        serie,
        attentions=prediction.attentions,
        save_directory=output_dir,
        gain=3,
    )

    print(f"Finished writing attention images to {output_dir}")

if __name__ == "__main__":
    main()