import gradio as gr
import base64
from PIL import Image
import io
import json

# 修改函式以確保它接收一個字典（這是 gradio_client 預設發送的格式）
def process(payload):
    try:
        # 如果客戶端已經傳送字典，直接使用 payload:
        data = payload 

        # decode base64 image
        img_bytes = base64.b64decode(data["image_b64"])
        img = Image.open(io.BytesIO(img_bytes))

        # This goes to Jetson
        reply = {
            "received": True,
            "robot_id": data.get("robot_id"),
            "size": img.size
        }

        # *** 關鍵修改：回傳一個包含圖片和 JSON 回覆的元組 (tuple) ***
        # Gradio 會自動將第一個值賦給第一個輸出元件 (gr.Image)
        # 第二個值賦給第二個輸出元件 (gr.JSON)
        return img, reply

    except Exception as e:
        # 發生錯誤時，確保回傳兩個值，其中圖片值為 None
        return None, {"error": str(e)}


demo = gr.Interface(
    fn=process,
    # 我們將輸入定義為 JSON，這允許後端接收字典格式
    inputs=gr.JSON(label="Input Payload (Dict format)"),
    outputs=[
        gr.Image(type="pil", label="Image Preview"), # 現在將接收 img 物件
        gr.JSON(label="Reply to Jetson")            # 現在將接收 reply 字典
    ],
    api_name="predict"
)

demo.launch()