anomalib/cli/install.py

"""Anomalib install subcommand code."""

# Copyright (C) 2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import logging

from pkg_resources import Requirement
from rich.console import Console
from rich.logging import RichHandler

from anomalib.cli.utils.installation import (
    get_requirements,
    get_torch_install_args,
    parse_requirements,
)

logger = logging.getLogger("pip")
logger.setLevel(logging.WARNING)  # setLevel: CRITICAL, ERROR, WARNING, INFO, DEBUG, NOTSET
console = Console()
handler = RichHandler(
    console=console,
    show_level=False,
    show_path=False,
)
logger.addHandler(handler)


def anomalib_install(option: str = "full", verbose: bool = False) -> int:
    """Install Anomalib requirements.

    Args:
        option (str | None): Optional-dependency to install requirements for.
        verbose (bool): Set pip logger level to INFO

    Raises:
        ValueError: When the task is not supported.

    Returns:
        int: Status code of the pip install command.
    """
    from pip._internal.commands import create_command

    requirements_dict = get_requirements("anomalib")

    requirements = []
    if option == "full":
        for extra in requirements_dict:
            requirements.extend(requirements_dict[extra])
    elif option in requirements_dict:
        requirements.extend(requirements_dict[option])
    elif option is not None:
        requirements.append(Requirement.parse(option))

    # Parse requirements into torch and other requirements.
    # This is done to parse the correct version of torch (cpu/cuda).
    torch_requirement, other_requirements = parse_requirements(requirements, skip_torch=option not in ("full", "core"))

    # Get install args for torch to install it from a specific index-url
    install_args: list[str] = []
    torch_install_args = []
    if option in ("full", "core") and torch_requirement is not None:
        torch_install_args = get_torch_install_args(torch_requirement)

    # Combine torch and other requirements.
    install_args = other_requirements + torch_install_args

    # Install requirements.
    with console.status("[bold green]Installing packages...  This may take a few minutes.\n") as status:
        if verbose:
            logger.setLevel(logging.INFO)
            status.stop()
        console.log(f"Installation list: [yellow]{install_args}[/yellow]")
        status_code = create_command("install").main(install_args)
        if status_code == 0:
            console.log(f"Installation Complete: {install_args}")

    if status_code == 0:
        console.print("Anomalib Installation [bold green]Complete.[/bold green]")

    return status_code