import numpy as np
from PIL import Image
import gradio as gr

import create_model as cm

print("Loading model and tokenizer...")
model_tokenizer = cm.create_model()
print("Model ready.")


def predict(image_1, image_2):
    if image_1 is None:
        return "Please upload at least the first X-ray image."

    img1 = np.array(Image.fromarray(image_1).convert("RGB")) / 255
    if image_2 is None:
        img2 = img1
    else:
        img2 = np.array(Image.fromarray(image_2).convert("RGB")) / 255

    caption = cm.function1([img1], [img2], model_tokenizer)
    return caption[0]


examples = [
    ["test_images/1/CXR54_IM-2145-1001.png", "test_images/1/CXR54_IM-2145-1002.png"],
    ["test_images/2/images.jpg", None],
    ["test_images/3/CXR303_IM-1404-1001.png", None],
    ["test_images/4/CXR25_IM-1024-2001.png", None],
]

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(label="X-ray 1 (frontal view)"),
        gr.Image(label="X-ray 2 (lateral view, optional)"),
    ],
    outputs=gr.Textbox(label="Impression"),
    examples=examples,
    title="Chest X-ray Report Generator",
    description=(
        "Upload one or two chest X-rays (frontal view, and optionally a lateral view) "
        "of the same patient to generate the impression section of a radiology report."
    ),
)

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