scripts/inference.py

from scripts.constants import CLASS_NAMES


def infer(model, input_files):
    for p in input_files:
        print("Processing", p, "...")
        probs = model(p)
        result_str = "Result -- " + " / ".join(
            [f"{name}: {probs[i].item():.4f}" for i, name in enumerate(CLASS_NAMES)]
        )
        print(result_str + "\n")