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

def process(json_input):
    try:
        #data = json.loads(json_input)

        # 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
        }

        # return only the msg: (Jetson JSON)
        return reply

    except Exception as e:
        return None, {"error": str(e)}


demo = gr.Interface(
    fn=process,
    inputs=gr.JSON(label="Jetson JSON"),
    outputs=[
        gr.Image(type="pil", label="Image Preview"),
        gr.JSON(label="Reply to Jetson")
    ],
    api_name="predict"
)

demo.launch()