import gradio as gr
from PIL import Image
from llava_inference import LLaVAHelper

model = LLaVAHelper()

def answer_question(image, question):
    if image is None or question.strip() == "":
        return "Please upload an image and enter a question."
    return model.generate_answer(image, question)

demo = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.Image(type="pil", label="Upload Public Transport Signage"),
        gr.Textbox(label="Ask a question (e.g., 'When is the next train to London?')")
    ],
    outputs=gr.Textbox(label="Answer"),
    title="UK Public Transport Assistant",
    description="Upload an image of UK public transport signage (like train timetables or metro maps), and ask a question related to it. Powered by LLaVA-1.5.",
    examples=[
        ["assets/example.jpg", "Where is platform 3?"],
    ]
)

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