File size: 2,508 Bytes
ea8b89c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | """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()
|