import os
import torch
from PIL import Image
from transformers import VisionEncoderDecoderModel, AutoTokenizer

def test_unimernet(image_path):
    # 1. 画像の存在確認
    if not os.path.exists(image_path):
        print(f"Error: 以下の画像が見つかりません: {image_path}")
        return

    print("UniMERNet モデルをロード中... (初回は時間がかかります)")
    
    model_id = "wanderkid/unimernet_small"
    
    try:
        # モデルとプロセッサの取得
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = VisionEncoderDecoderModel.from_pretrained(model_id, trust_remote_code=True).to(device)
        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        model.eval()
        print(f"Using device: {device}")

        # 2. 画像の読み込み
        image = Image.open(image_path).convert("RGB")
        
        # 3. 認識実行
        print("画像を解析中...")
        with torch.no_grad():
            # UniMERNet のカスタムモデルクラスが predict メソッドを提供していることを想定
            latex_code = model.predict(image, tokenizer)
        
        print("\n--- 認識結果 (LaTeX) ---")
        print(latex_code)
        print("------------------------\n")
        
    except Exception as e:
        print(f"Error during inference: {e}")

if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1:
        test_unimernet(sys.argv[1])
    else:
        print("使い方: python predict.py <画像のパス>")