"""Classify a local image against candidate labels with MobileCLIP2.

Install: python -m pip install "open_clip_torch==3.3.0" "timm==1.0.29" torch pillow
Run: python mobileclip2.py image.png --labels 'a receipt' 'a chart' 'a photograph'
Uses the timm/OpenCLIP adaptation of Apple's weights, downloaded on first run.
Source: https://huggingface.co/timm/MobileCLIP2-S0-OpenCLIP
"""
import argparse
import json
from pathlib import Path

import open_clip
import torch
from PIL import Image
from timm.utils import reparameterize_model


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('image', type=Path)
    parser.add_argument('--model', choices=['S0', 'S2', 'B', 'S3', 'L-14', 'S4'], default='S0')
    parser.add_argument('--labels', nargs='+', required=True)
    args = parser.parse_args()
    if not args.image.is_file():
        parser.error('The image file does not exist.')
    if len(args.labels) < 2:
        parser.error('Provide at least two candidate labels.')
    name = 'MobileCLIP2-' + args.model
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model, _, preprocess = open_clip.create_model_and_transforms(name, pretrained='dfndr2b')
    model = reparameterize_model(model.eval()).to(device)
    tokenizer = open_clip.get_tokenizer(name)
    image = preprocess(Image.open(args.image).convert('RGB')).unsqueeze(0).to(device)
    text = tokenizer(args.labels).to(device)
    with torch.inference_mode():
        image_features = model.encode_image(image)
        text_features = model.encode_text(text)
        image_features /= image_features.norm(dim=-1, keepdim=True)
        text_features /= text_features.norm(dim=-1, keepdim=True)
        scores = (100.0 * image_features @ text_features.T).softmax(dim=-1)[0].cpu().tolist()
    results = sorted(zip(args.labels, scores), key=lambda item: item[1], reverse=True)
    print(json.dumps({'model': name, 'scores': [{'label': label, 'score': score} for label, score in results],
        'note': 'Relative scores within the supplied labels, not calibrated probabilities.'}, indent=2))


if __name__ == '__main__':
    main()