app.py

# app.py (Hugging Face Spaces用)

import streamlit as st
from transformers import pipeline
from PIL import Image
import time

# モデルのパス設定
# 🚨🚨 Spacesリポジトリにアップロードしたモデルフォルダの相対パス 🚨🚨
# 例: vit-kuzushiji49-final-tpu-xla-speed
MODEL_NAME = "./vit-kuzushiji49-final-tpu-xla-speed" 

# 2. モデルの初期化（@st.cache_resource でキャッシュし、ロードを高速化）
# device=-1 はCPUを意味します (Spacesの無料枠はCPU)
@st.cache_resource
def load_classifier():
    st.info("モデルをロード中... 数分かかることがあります。")
    try:
        classifier = pipeline(
            "image-classification", 
            model=MODEL_NAME,
            device=-1 # CPUでの推論を指定
        )
        return classifier
    except Exception as e:
        st.error(f"モデルのロードに失敗しました。パスを確認してください: {e}")
        return None

# --- アプリケーションのメインロジック ---
st.set_page_config(page_title="崩し字分類", layout="wide")
st.title("🏯 崩し字 49 分類アプリ (ViT)")
st.markdown("Hugging Face SpacesのCPU上で動作しています。")

classifier = load_classifier()

if classifier:
    uploaded_file = st.file_uploader("崩し字の画像をアップロードしてください (JPEG/PNG)", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        # 画像をPIL Imageとして読み込む
        image = Image.open(uploaded_file).convert('RGB')
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.image(image, caption='アップロードされた画像', use_column_width=True)
            
        with col2:
            st.subheader("分類結果")
            
            start_time = time.time()
            
            # 推論の実行
            with st.spinner('分類中...'):
                results = classifier(image)

            end_time = time.time()
            
            # 結果の表示
            if results:
                st.success(f"推定される文字: **{results[0]['label']}** (確率: {results[0]['score']:.4f})")
                st.write(f"推論時間: {end_time - start_time:.3f}秒 (CPU)")
                
                # 詳細な結果をテーブルで表示
                st.dataframe({
                    "順位": list(range(1, len(results) + 1)),
                    "ラベル (漢字)": [res['label'] for res in results[:5]], 
                    "確率": [f"{res['score']:.4f}" for res in results[:5]]
                }, hide_index=True)