import gradio
import gradio as gr
import os

from SwinIR.infer import SwinIRDemo
from sam_diffsr.infer import sam_diffsr_demo


def image_infer(img_PIL):
    sam_diffsr_img = sam_diffsr_infer.infer(img_PIL)
    swin_ir_img = swin_ir_infer.infer(img_PIL)
    return sam_diffsr_img, swin_ir_img


if __name__ == "__main__":
    sam_diffsr_infer = sam_diffsr_demo()
    swin_ir_infer = SwinIRDemo()
    
    root_path = os.path.dirname(__file__)
    cheetah = os.path.join(root_path, "images/04011.png")
    
    demo = gr.Interface(image_infer, gr.Image(type="pil", value=cheetah),
                        [
                                gradio.Image(label='SAM-DiffSR', show_label=True),
                                gradio.Image(label='SwinIR', show_label=True)
                        ],
                        # flagging_options=["blurry", "incorrect", "other"],
                        examples=[
                                os.path.join(root_path, "images/04011.png"),
                                os.path.join(root_path, "images/04033.png"),
                                os.path.join(root_path, "images/04064.png"),
                                os.path.join(root_path, "images/04146.png"),
                                # os.path.join(root_path, "images/10091.png"),
                                os.path.join(root_path, "images/0801x4.png"),
                                os.path.join(root_path, "images/0804x4.png"),
                                os.path.join(root_path, "images/0809x4.png"),
                        ]
                        )
    
    demo.launch()