import gradio as gr
import base64
from openai import OpenAI
import pandas as pd
import json
import io
import cv2
def estimate_calories(api_key, image):
    client = OpenAI(api_key=api_key)

    # 画像をbase64エンコードする
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    _, buffer = cv2.imencode(".jpg", image_rgb)
    base64_image = base64.b64encode(buffer).decode("utf-8")

    # GPT-4にカロリーを推定させる (JSON-MODE)
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": """あなたは画像からカロリーを推定する優秀な栄養士です。提供された食事の画像を分析し、カロリーを推定してください。""",
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpg;base64,{base64_image}",
                            "detail": "high",
                        },
                    },
                    {
                        "type": "text",
                        "text": """上記の食事の画像からカロリーを推定してください。料理名は日本語で回答してください。
## JSON Schema

```
{
    "type": "object",
    "properties": {
        "step_by_step_estimation": {"type": "string", "description": "料理とカロリーについての推論"},
        "foods": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "name": {"type": "string", "description": "料理名"},
                    "calorie": {"type": "string", "description": "カロリー"},
                },
                "required": ["name", "calorie"],
            },
        }
    },
    "required": ["step_by_step_estimation", "foods"],
}
```
""",
                    },
                ],
            },
        ],
        temperature=1,
        response_format={"type": "json_object"},
    )
    content_json = json.loads(response.choices[0].message.content)
    step_by_step_estimation = content_json["step_by_step_estimation"]
    foods = content_json["foods"]

    # foodsをデータフレームに変換する
    df = pd.DataFrame(foods)
    return step_by_step_estimation, df

demo = gr.Interface(
    fn=estimate_calories,
    inputs=[gr.Textbox(label="OpenAI API Key"), gr.Image(label="Upload Food Image")],
    outputs=[
        gr.Textbox(label="段階的な推論"),
        gr.DataFrame(headers=["name", "calorie"], label="推定カロリー"),
    ],
    title="Calorie Predictor",
    description="食事の画像をアップロードしOpenAIのAPIキーを入力すると、カロリーが推定されます。API使用料にご注意ください。",
)

if __name__ == "__main__":
    demo.launch()