diff --git a/brainles_aurora/inferer/inferer.py b/brainles_aurora/inferer/inferer.py index 5546bd5..eef4d44 100644 --- a/brainles_aurora/inferer/inferer.py +++ b/brainles_aurora/inferer/inferer.py @@ -17,6 +17,7 @@ from brainles_aurora.inferer.data import DataHandler from brainles_aurora.inferer.model import ModelHandler from brainles_aurora.utils import remove_path_suffixes +from brainles_aurora.utils.console_decorators import citation_reminder logger = logging.getLogger(__name__) @@ -148,6 +149,7 @@ def _configure_device(self) -> torch.device: return device + @citation_reminder def infer( self, t1: str | Path | np.ndarray | None = None, diff --git a/brainles_aurora/utils/console_decorators.py b/brainles_aurora/utils/console_decorators.py new file mode 100644 index 0000000..dc60b7b --- /dev/null +++ b/brainles_aurora/utils/console_decorators.py @@ -0,0 +1,22 @@ +from rich.console import Console + +CITATION_LINK = "https://github.com/BrainLesion/AURORA#citation" + + +def citation_reminder(func): + def wrapper(*args, **kwargs): + console = Console() + console.rule("Thank you for using [bold]AURORA[/bold]") + console.print( + f"Please support our development by citing the papers listed here:", + justify="center", + ) + console.print( + f"{CITATION_LINK} -- Thank you!", + justify="center", + ) + console.rule() + console.line() + func(*args, **kwargs) + + return wrapper diff --git a/pyproject.toml b/pyproject.toml index 06261b7..99c418c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ numpy = ">=1.23.0" tqdm = "^4.64.1" path = "^16.2.0" requests =">=2.0.0" +rich = ">=13.0.0" [tool.poetry.dev-dependencies] pytest = "^6.2"