File size: 1,124 Bytes
9524435 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | """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()
|