Spaces:
Running
Running
| import torch | |
| import numpy as np | |
| import gradio as gr | |
| from PIL import Image | |
| from models import dehazeformer | |
| def infer(raw_image): | |
| network = dehazeformer() | |
| network.load_state_dict(torch.load('./saved_models/dehazeformer.pth', map_location=torch.device('cpu'))['state_dict']) | |
| # torch.save({'state_dict': network.state_dict()}, './saved_models/dehazeformer.pth') | |
| network.eval() | |
| image = np.array(raw_image, np.float32) / 255. * 2 - 1 | |
| image = torch.from_numpy(image) | |
| image = image.permute((2, 0, 1)).unsqueeze(0) | |
| with torch.no_grad(): | |
| output = network(image).clamp_(-1, 1)[0] * 0.5 + 0.5 | |
| output = output.permute((1, 2, 0)) | |
| output = np.array(output, np.float32) | |
| output = np.round(output * 255.0) | |
| output = Image.fromarray(output.astype(np.uint8)) | |
| return output | |
| title = "DehazeFormer" | |
| description = f"We use a mixed dataset to train the model, allowing the trained model to work better on real hazy images. To allow the model to process high-resolution images more efficiently and effectively, we extend it to the [MCT](https://github.com/IDKiro/MCT) variant." | |
| examples = [ | |
| ["examples/1.jpg"], | |
| ["examples/2.jpg"], | |
| ["examples/3.jpg"], | |
| ["examples/4.jpg"], | |
| ["examples/5.jpg"], | |
| ["examples/6.jpg"] | |
| ] | |
| iface = gr.Interface( | |
| infer, | |
| inputs="image", outputs="image", | |
| title=title, | |
| description=description, | |
| allow_flagging='never', | |
| examples=examples, | |
| ) | |
| iface.launch() |