import gradio as gr
import cv2
import numpy as np
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import load_model

# TensorFlow.jsモデルのURL
URL = "https://teachablemachine.withgoogle.com/models/ZPfAhDYCh/"

# モデルのロード
model_url = URL + "model.json"
metadata_url = URL + "metadata.json"
model = load_model(model_url)

# クラスの数を取得
max_predictions = model.layers[-1].output_shape[1]

# カメラの初期化
def init():
    global cap
    cap = cv2.VideoCapture(0)  # カメラのデフォルトデバイスを使用

# 画像を予測
def predict():
    global cap
    ret, frame = cap.read()  # カメラからフレームをキャプチャ
    if not ret:
        print("Failed to capture image from camera.")
        return

    frame = cv2.resize(frame, (200, 200))  # 画像サイズをモデルの入力サイズに変更
    img_array = image.img_to_array(frame)
    img_array = np.expand_dims(img_array, axis=0)
    img_array /= 255.0  # 画像データの正規化

    # 予測
    prediction = model.predict(img_array)
    
    # 結果の表示
    result = {}
    for i in range(max_predictions):
        result[f"Class {i}"] = prediction[0][i]

    return result

# インターフェースの作成
iface = gr.Interface(fn=predict, live=True, capture_session=True)
iface.launch()