import gradio as gr
import cv2
import os

def restore(image_missing, image_inference, channel):
    if len(image_missing.shape) != 3 or len(image_inference.shape) != 3:
        return None #[None, "Error: Please provide RGB images!"]
    c = { "R": 0, "G": 1, "B": 2 }[channel] # cv2 uses BGR order
    image_missing[:,:,c] = image_inference[:,:,c]
    return image_missing

examples=[
        [os.path.join(os.path.dirname(__file__), "missing/bird_r.webp"  ), os.path.join(os.path.dirname(__file__), "inference/bird_r.webp"  ), "R"],
        [os.path.join(os.path.dirname(__file__), "missing/bird_g.webp"  ), os.path.join(os.path.dirname(__file__), "inference/bird_g.webp"  ), "G"],
        [os.path.join(os.path.dirname(__file__), "missing/bird_b.webp"  ), os.path.join(os.path.dirname(__file__), "inference/bird_b.webp"  ), "B"],
        [os.path.join(os.path.dirname(__file__), "missing/dog2_r.webp"  ), os.path.join(os.path.dirname(__file__), "inference/dog2_r.webp"  ), "R"],
        [os.path.join(os.path.dirname(__file__), "missing/dog2_g.webp"  ), os.path.join(os.path.dirname(__file__), "inference/dog2_g.webp"  ), "G"],
        [os.path.join(os.path.dirname(__file__), "missing/dog2_b.webp"  ), os.path.join(os.path.dirname(__file__), "inference/dog2_b.webp"  ), "B"],
        [os.path.join(os.path.dirname(__file__), "missing/house2_r.webp"), os.path.join(os.path.dirname(__file__), "inference/house2_r.webp"), "R"],
        [os.path.join(os.path.dirname(__file__), "missing/house2_g.webp"), os.path.join(os.path.dirname(__file__), "inference/house2_g.webp"), "G"],
        [os.path.join(os.path.dirname(__file__), "missing/house2_b.webp"), os.path.join(os.path.dirname(__file__), "inference/house2_b.webp"), "B"],
    ]

app = gr.Interface(
    fn=restore,
    inputs=[
            gr.Image("missing/dog2_r.webp", type="numpy", label="Missing channel"), 
            gr.Image("inference/dog2_r.webp", type="numpy", label="Inference"), 
            gr.Radio(["R", "G", "B"], label="Restore channel"),
        ],    
    outputs="image",
    examples=examples,
    title="Restore missing RGB channel",
    description="Restore missing channel of a RGB image by using a ControlNet to guide image generation of Stable Diffusion to infer missing channel from the other two channels https://huggingface.co/GeroldMeisinger/control-channels . Examples are missing one channel each (use `decompose.py`) and inference examples are cherry-picked from generation with ControlNet. The submit button replaces the missing channel from the inference. To generate your own inferences you have to set up Stable Diffusion with ControlNet 'channels-rgb'. Note: While the inference images look fine on their own and you might as well use them, we are only interested in one of their channels.")

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