predict.py

#!/usr/bin/env python3
"""
CLI Predictor for Captcha Recognition Model.
Developed by Confia Company.

Usage:
    python predict.py /path/to/captcha.png
    python predict.py /path/to/captcha.png --weights model.pt
    echo /path/to/captcha.png | python predict.py
"""

import argparse
import os
import sys
import torch
from modeling_captcha import load_model, predict_captcha, DEFAULT_ALPHABET

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_WEIGHTS = os.path.join(HERE, "model.pt")


def main():
    parser = argparse.ArgumentParser(description="Predict 5-character captcha answer")
    parser.add_argument("image", nargs="?", help="Path to captcha image file (PNG/JPG)")
    parser.add_argument("--weights", default=DEFAULT_WEIGHTS, help=f"Path to model weights (default: {DEFAULT_WEIGHTS})")
    args = parser.parse_args()

    img_path = args.image
    if not img_path:
        img_path = sys.stdin.read().strip()
    if not img_path:
        parser.error("No image path provided.")

    if not os.path.isfile(img_path):
        print(f"Error: Image file '{img_path}' not found.", file=sys.stderr)
        sys.exit(1)

    weights = args.weights
    if not os.path.isfile(weights):
        # Fallback to best.pt or pytorch_model.bin if model.pt is missing
        for alt in ["best.pt", "pytorch_model.bin"]:
            alt_path = os.path.join(HERE, alt)
            if os.path.isfile(alt_path):
                weights = alt_path
                break

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = load_model(weights, alphabet=DEFAULT_ALPHABET, device=device)
    result = predict_captcha(model, img_path, alphabet=DEFAULT_ALPHABET, device=device)
    print(result)


if __name__ == "__main__":
    main()