from ..media import copy_frames
from ..models.ben2 import infer_foreground
from .common import input_frames


def run(context, parameters):
    context.progress(0.02, "Preparing BEN2 API input")
    context.directory.mkdir(parents=True, exist_ok=True)
    model_manager = context.resource("model_manager")
    model_directory = model_manager.directory(parameters["model"])
    with input_frames(context, parameters) as images:
        for index, image in enumerate(images):
            context.check_cancelled()
            foreground, _load_ms, _inference_ms = infer_foreground(
                image,
                model_directory,
                parameters["device"],
                model_cache=model_manager,
                cancel_check=context.check_cancelled,
            )
            suffix = "" if len(images) == 1 else f".{index + 1:04d}"
            foreground.save(context.directory / f"foreground{suffix}.png")
        copy_frames(images, context.directory)
    context.check_cancelled()
    context.progress(1.0, "Saved raw BEN2 API foreground")
    return {"directory": "."}