import gradio as gr
import os
from inference import Pipeline

# Initialize the pipeline
# This will load the model onto the GPU (if available)
pipeline = Pipeline(model_path="models for IML/APSC-Net.pth")

def predict(image):
    """
    Function to be called by the Gradio interface.
    Takes a PIL image and returns the mask.
    """
    # The pipeline returns a dictionary with a base64 encoded image string
    # We don't need to decode it as Gradio's Image component can handle it.
    result = pipeline(image)
    # The output format for gr.Image can be a PIL Image, numpy array, or file path.
    # We can directly return the base64 string with a data URI scheme.
    return f"data:image/png;base64,{result['image']}"

# --- Gradio Interface ---
title = "APSC-Net: Image Manipulation Localization"
description = """
Official Gradio Demo for **APSC-Net**, from the paper *"Towards Modern Image Manipulation Localization: A Large-Scale Dataset and Novel Methods"*.
Upload an image to detect manipulated regions. The model will output a binary mask where white indicates tampered regions.
\nFor more details, visit the [project repository](https://github.com/qcf-568/MIML).
"""

# Example images from the 'models for IML' directory (assuming you have a test folder)
# Add your own example images here
example_paths = [
    os.path.join(os.path.dirname(__file__), "models/for_IML/test_data/CASIA1/imgs/Sp_D_CRN_A_cha0001_cha0001_0101_1.jpg"),
    os.path.join(os.path.dirname(__file__), "models/for_IML/test_data/IMD20/imgs/0000_000_02.png"),
]
# Filter out examples that don't exist
examples = [path for path in example_paths if os.path.exists(path)]


iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Input Image"),
    outputs=gr.Image(type="pil", label="Predicted Tamper Mask"),
    title=title,
    description=description,
    examples=examples,
    cache_examples=True,
    allow_flagging="never",
    article="<p style='text-align: center'>Paper: <a href='https://openaccess.thecvf.com/content/CVPR2024/papers/Qu_Towards_Modern_Image_Manipulation_Localization_A_Large-Scale_Dataset_and_Novel_CVPR_2024_paper.pdf' target='_blank'>Towards Modern Image Manipulation Localization: A Large-Scale Dataset and Novel Methods</a></p>"
)

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