| """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() |
|
|
