app.py

import torch
import gradio as gr
import numpy as np
import torch.nn.functional as F

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

model_restoration = Histoformer()
checkpoint = torch.load("Allweather/pretrained_models/net_g_real.pth")
model_restoration.load_state_dict(checkpoint['params'])
model_restoration.cuda()
model_restoration = nn.DataParallel(model_restoration)

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

    # 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')
    return input_

def predict(input_img):
    prediction = model_restoration(preprocess(input_img))
    return input_img, prediction 

gradio_app = gr.Interface(
    predict,
    inputs=gr.Image(label="Upload images with adverse weather degradations", sources=['upload', 'webcam'], type="pil"),
    outputs=gr.Image(label="Processed Image"),
    title="Image Restoration for All-in-one Adverse Weather",
)

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