"""Classify a generated image, or pass --image to classify your own image's color."""
import argparse
import json
from pathlib import Path

from PIL import Image
from jevlike import Choice, Noul
from jevlike.mm import MultimodalSystemOne


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--model-dir', type=Path, default=Path(__file__).resolve().parents[1])
    parser.add_argument('--device', default='cuda', choices=['cuda', 'cpu'])
    parser.add_argument('--image', type=Path)
    args = parser.parse_args()
    model = MultimodalSystemOne.load(str(args.model_dir), device=args.device,
                                    vision=str(args.model_dir / 'vision'))
    questions = {
        'color': Choice('Which color dominates the picture?', ['red', 'green', 'blue']),
        'red': Noul('The picture is predominantly red.'),
    }
    picture = str(args.image) if args.image else Image.new('RGB', (224, 224), 'red')
    result = model.predict(None, questions, image=picture)
    print(json.dumps(result, indent=2, allow_nan=False))


if __name__ == '__main__':
    main()