Handheld / app.py
Rausda6's picture
Add minimal Gradio burst reconstruction frontend
ea8b89c verified
Raw History Blame Contribute Delete
2.51 kB
"""Minimal Gradio frontend for raw burst reconstruction."""
from pathlib import Path
import shutil
import tempfile
import spaces
import gradio as gr
import numpy as np
@spaces.GPU(duration=120)
def reconstruct(files, scale):
# Import CUDA modules only after a GPU has been allocated.
from handheld_super_resolution.config import Config
from handheld_super_resolution.monte_carlo import generate_noise_lut
from handheld_super_resolution.noise_lut import save_noise_lut
from handheld_super_resolution.utils_dng import read_dng_noise_profile
from handheld_super_resolution.super_resolution import process
with tempfile.TemporaryDirectory(prefix="handheld-") as directory:
burst = Path(directory)
for index, source in enumerate(sorted(files, key=lambda p: Path(p).name)):
shutil.copyfile(source, burst / f"{index:04d}.dng")
alpha, beta = read_dng_noise_profile(burst / "0000.dng")
lut = generate_noise_lut(alpha, beta, show_progress=False)
lut_path = burst / "noise.npz"
save_noise_lut(lut_path, lut, trials=10_000_000, seed=0)
config = Config(scale=float(scale))
config.noise_model.lut_path = lut_path
config.robustness.save_mask = False
result, _ = process(burst, config)
return np.round(np.clip(np.nan_to_num(result), 0, 1) * 255).astype(np.uint8)
def run(files, scale):
if not files or len(files) < 2:
raise gr.Error("Upload at least two DNG images from the same burst.")
if any(Path(path).suffix.lower() != ".dng" for path in files):
raise gr.Error("This pipeline accepts raw DNG images.")
try:
return reconstruct(files, scale)
except Exception as error:
raise gr.Error(f"Reconstruction failed: {error}") from error
with gr.Blocks(delete_cache=(3600, 3600)) as demo:
gr.Markdown("# Handheld Super-Resolution\nUpload raw DNG images from one burst. "
"The first image alphabetically is the reference. "
"Noise calibration is automatic; the first run may take longer.")
images = gr.File(label="DNG burst", file_count="multiple", file_types=[".dng"], type="filepath")
scale = gr.Radio([1, 2], value=2, label="Output scale")
button = gr.Button("Reconstruct", variant="primary")
result = gr.Image(label="Result", format="png")
button.click(run, inputs=[images, scale], outputs=result, concurrency_limit=1)
if __name__ == "__main__":
demo.queue(max_size=8).launch()