import torch
import gradio as gr
import numpy as np
import torch.nn.functional as F
from skimage import img_as_ubyte

from Allweather.util import load_img, save_img
from basicsr.models.archs.histoformer_arch import Histoformer

model_restoration = Histoformer.from_pretrained("sunsean/Histoformer-real")

model_restoration.eval()

factor = 8
def predict(input_img):
    img = np.float32(load_img(input_img))/255.
    img = torch.from_numpy(img).permute(2,0,1)
    input_ = img.unsqueeze(0)

    # Padding in case images are not multiples of 8
    h,w = input_.shape[2], input_.shape[3]
    H,W = ((h+factor)//factor)*factor, ((w+factor)//factor)*factor
    padh = H-h if h%factor!=0 else 0
    padw = W-w if w%factor!=0 else 0
    input_ = F.pad(input_, (0,padw,0,padh), 'reflect')
    
    restored = model_restoration(input_)
    output_path = "restored.png"
    restored = restored[:,:,:h,:w]
    restored = torch.clamp(restored,0,1).detach().permute(0, 2, 3, 1).squeeze(0).numpy()

    save_img(output_path, img_as_ubyte(restored))

example_images = [
    "examples/example.jpg",
]
gradio_app = gr.Interface(
    predict,
    inputs=gr.Image(label="Upload images with adverse weather degradations", type="filepath"),
    outputs=[
        gr.Image(type="filepath", label="Restored image", height=768, width=768),
        gr.Textbox(label="Error Message")
    ],
    title="Histoformer: All-in-one Image Restoration under Adverse Weather Conditions",
    description="[Histoformer](https://huggingface.co/sunsean/Histoformer/) is a image restoration model for all-in-one adverse weather.",
    examples=example_images
)

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