|
import os |
|
|
|
import gradio as gr |
|
import numpy as np |
|
import torch |
|
from torchvision.utils import make_grid |
|
from huggingface_hub import snapshot_download |
|
from zero_dce import enhance_net_nopool |
|
|
|
os.system("pip freeze") |
|
|
|
REPO_ID = "leonelhs/lowlight" |
|
MODEL_NAME = "Epoch99.pth" |
|
|
|
model = enhance_net_nopool().cpu() |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
snapshot_folder = snapshot_download(repo_id=REPO_ID) |
|
model_path = os.path.join(snapshot_folder, MODEL_NAME) |
|
state = torch.load(model_path, map_location=device) |
|
model.load_state_dict(state) |
|
|
|
|
|
def tensor_to_ndarray(tensor, nrow=1, padding=0, normalize=True): |
|
grid = make_grid(tensor, nrow, padding, normalize) |
|
return grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() |
|
|
|
|
|
def predict(image): |
|
image = (np.asarray(image) / 255.0) |
|
image = torch.from_numpy(image).float() |
|
image = image.permute(2, 0, 1) |
|
image = image.cpu().unsqueeze(0) |
|
_, enhanced_image, _ = model(image) |
|
return tensor_to_ndarray(enhanced_image, nrow=8, padding=2, normalize=False) |
|
|
|
|
|
title = "Zero-DCE" |
|
description = r""" |
|
## Low-Light Image Enhancement using Zero-DCE |
|
|
|
The model improves the quality of images that have poor contrast, low brightness, and suboptimal exposure. |
|
|
|
This is an implementation of <a href='https://github.com/Li-Chongyi/Zero-DCE' target='_blank'>Zero-DCE</a>. |
|
|
|
It has no any particular purpose than start research on AI models. |
|
|
|
""" |
|
|
|
article = r""" |
|
Questions, doubts, comments, please email 📧 `leonelhs@gmail.com` |
|
|
|
This demo is running on a CPU, if you like this project please make us a donation to run on a GPU or just give us a <a href='https://github.com/leonelhs/Zero-DCE' target='_blank'>Github ⭐</a> |
|
|
|
<a href="https://www.buymeacoffee.com/leonelhs"><img src="https://img.buymeacoffee.com/button-api/?text=Buy me a coffee&emoji=&slug=leonelhs&button_colour=FFDD00&font_colour=000000&font_family=Cookie&outline_colour=000000&coffee_colour=ffffff" /></a> |
|
|
|
<center><img src='https://visitor-badge.glitch.me/badge?page_id=zerodce.visitor-badge' alt='visitor badge'></center> |
|
""" |
|
|
|
demo = gr.Interface( |
|
predict, [ |
|
gr.Image(type="pil", label="Image low light"), |
|
], [ |
|
gr.Image(type="numpy", label="Image enhanced") |
|
], |
|
title=title, |
|
description=description, |
|
article=article) |
|
|
|
demo.queue().launch() |
|
|