import gradio as gr import numpy as np from PIL import Image from pystackreg import StackReg import imageio.v2 as iio import tifffile import tempfile import urllib.request import os from core.utils import ( get_sr_mode, normalize_stack, upscale, load_stack, _is_demo_url, _demo_path_for_url, _start_cleaner, citation_markdown, documentation_markdown ) _start_cleaner() # Globals original_frames, aligned_frames = [], [] ref_frames, mov_frames, reg_frames = [], [], [] custom_stack, reg_result = [], None # Reset functions def reset_intra_stack(): global original_frames, aligned_frames original_frames, aligned_frames = [], [] return [None, gr.update(value=0, minimum=0, maximum=0), None, gr.update(value=0, minimum=0, maximum=0), None, gr.update(value=0, minimum=0, maximum=0), None, gr.update(value=0, minimum=0, maximum=0), None] def reset_reference_based(): global ref_frames, mov_frames, reg_frames ref_frames, mov_frames, reg_frames = [], [], [] return [None, None, None, gr.update(value=0, minimum=0, maximum=0), None, gr.update(value=0, minimum=0, maximum=0), None] def reset_frame_to_frame(): global custom_stack, reg_result custom_stack, reg_result = [], None return [None, gr.update(value=0, minimum=0, maximum=0), gr.update(value=0, minimum=0, maximum=0), None, None] # Registration logic def intra_stack_align(f, ref_idx, ext_file, ext_idx, mode): global original_frames, aligned_frames stack = load_stack(f) original_frames = [Image.fromarray(fr) for fr in stack] sr = StackReg(get_sr_mode(mode)) ref = load_stack(ext_file)[ext_idx] if ext_file else stack[ref_idx] aligned = [sr.register_transform(ref, fr) for fr in stack] aligned = normalize_stack(np.stack(aligned)) aligned_frames = [upscale(Image.fromarray(fr)) for fr in aligned] path = tempfile.NamedTemporaryFile(suffix=".tif", delete=False).name tifffile.imwrite(path, aligned, photometric='minisblack') return ( original_frames[0], gr.update(value=0, maximum=len(original_frames)-1), aligned_frames[0], gr.update(value=0, maximum=len(aligned_frames)-1), path ) def reference_align(ref_file, mov_file, mode): global ref_frames, mov_frames, reg_frames ref_stack = load_stack(ref_file) mov_stack = load_stack(mov_file) ref_frames = [Image.fromarray(f) for f in ref_stack] mov_frames = [Image.fromarray(f) for f in mov_stack] sr = StackReg(get_sr_mode(mode)) aligned = [sr.register_transform(ref_stack[0], f) for f in mov_stack] aligned = normalize_stack(np.stack(aligned)) reg_frames = [upscale(Image.fromarray(f)) for f in aligned] path = tempfile.NamedTemporaryFile(suffix=".tif", delete=False).name tifffile.imwrite(path, aligned, photometric='minisblack') return ref_frames[0], gr.update(value=0, maximum=len(ref_frames)-1), \ reg_frames[0], gr.update(value=0, maximum=len(reg_frames)-1), path def frame_to_frame_align(file, ref_idx, mov_idx, mode): global custom_stack, reg_result stack = load_stack(file) custom_stack = [Image.fromarray(f) for f in stack] sr = StackReg(get_sr_mode(mode)) aligned = sr.register_transform(stack[ref_idx], stack[mov_idx]) aligned = normalize_stack(np.stack([aligned]))[0] reg_result = Image.fromarray(aligned) path = tempfile.NamedTemporaryFile(suffix=".tif", delete=False).name tifffile.imwrite(path, aligned[np.newaxis, ...], photometric='minisblack') return reg_result, path # URL loading def load_from_url(request: gr.Request): params = request.query_params if "file_url" in params: try: url = params["file_url"] if _is_demo_url(url): tmp_path = _demo_path_for_url(url) if not os.path.exists(tmp_path): urllib.request.urlretrieve(url, tmp_path) else: tmp_path = tempfile.mktemp(suffix=".tif") # goes to WORK_DIR urllib.request.urlretrieve(url, tmp_path) return [gr.update(value=tmp_path)] except Exception as e: print(f"[URL load failed] {e}") return [None] # Interface with gr.Blocks() as demo: gr.Markdown("# 🧠 Pystackreg Web Application") gr.Markdown(citation_markdown) with gr.Accordion("📘 How to Use This App (Click to Expand)", open=False): gr.Markdown(documentation_markdown) with gr.Tab("📚 Reference-Based Alignment"): with gr.Row(): file_input = gr.File(label="Upload Stack to Register") use_ext_ref = gr.Checkbox(label="Use external reference stack") gr.Examples( examples=[ ["https://github.com/glichtner/pystackreg/raw/master/examples/data/pc12-unreg.tif"] ], inputs=[file_input], label="Try with Example TIFF" ) with gr.Row(): ref_slider = gr.Slider(label="Reference Frame (from uploaded stack)", minimum=0, maximum=0, value=0, step=1, visible=True) ext_ref_file = gr.File(label="Upload External Reference Stack (.tif)", visible=False) ext_ref_slider = gr.Slider(label="Reference Frame (from external stack)", minimum=0, maximum=0, value=0, step=1, visible=False) use_ext_ref.change( lambda v: ( gr.update(visible=not v), gr.update(visible=v), gr.update(visible=v) ), use_ext_ref, [ref_slider, ext_ref_file, ext_ref_slider] ) ext_ref_file.change( lambda f: gr.update(value=0, maximum=len(iio.mimread(f)) - 1) if f else gr.update(value=0, maximum=0), ext_ref_file, ext_ref_slider ) with gr.Row(): show_adv = gr.Checkbox(label="Show Advanced Settings", value=False) mode_dropdown = gr.Dropdown(["TRANSLATION", "RIGID_BODY", "SCALED_ROTATION", "AFFINE", "BILINEAR"], value="RIGID_BODY", visible=False, label="Transformation Mode") show_adv.change(lambda v: gr.update(visible=v), show_adv, mode_dropdown) run_btn = gr.Button("▶️ Align Stack") with gr.Row(): original_image = gr.Image(label="Original Frame") aligned_image = gr.Image(label="Aligned Frame") with gr.Row(): original_slider = gr.Slider(label="Browse Original Stack", minimum=0, maximum=0, value=0, step=1) aligned_slider = gr.Slider(label="Browse Aligned Stack", minimum=0, maximum=0, value=0, step=1) download = gr.File(label="Download") file_input.change( lambda f: gr.update(value=0, maximum=len(iio.mimread(f)) - 1) if f else gr.update(value=0, maximum=0), file_input, ref_slider ) run_btn.click( intra_stack_align, [file_input, ref_slider, ext_ref_file, ext_ref_slider, mode_dropdown], [original_image, original_slider, aligned_image, aligned_slider, download] ) original_slider.change(lambda i: original_frames[i] if 0 <= i < len(original_frames) else None, original_slider, original_image) aligned_slider.change(lambda i: aligned_frames[i] if 0 <= i < len(aligned_frames) else None, aligned_slider, aligned_image) gr.Button("🔄 Reset Tab").click( reset_intra_stack, outputs=[ file_input, ref_slider, ext_ref_file, ext_ref_slider, original_image, original_slider, aligned_image, aligned_slider, download ] ) with gr.Tab("🎯 Stack-Based Alignment"): with gr.Row(): ref_input = gr.File(label="Reference Stack") mov_input = gr.File(label="Moving Stack") gr.Examples( examples=[ [ "https://github.com/glichtner/pystackreg/raw/master/examples/data/pc12-unreg.tif", "https://github.com/glichtner/pystackreg/raw/master/examples/data/pc12-reg-translation.tif" ] ], inputs=[ref_input, mov_input], label="Try with Example Stacks" ) with gr.Row(): show_adv_ref = gr.Checkbox(label="Show Advanced Settings", value=False) mode_dropdown_ref = gr.Dropdown(["TRANSLATION", "RIGID_BODY", "SCALED_ROTATION", "AFFINE", "BILINEAR"], value="RIGID_BODY", visible=False, label="Transformation Mode") show_adv_ref.change(lambda v: gr.update(visible=v), show_adv_ref, mode_dropdown_ref) ref_btn = gr.Button("▶️ Register") with gr.Row(): ref_image = gr.Image(label="Reference Frame") reg_image = gr.Image(label="Registered Frame") with gr.Row(): ref_slider = gr.Slider(label="Browse Ref", minimum=0, maximum=0, value=0, step=1) reg_slider = gr.Slider(label="Browse Reg", minimum=0, maximum=0, value=0, step=1) download_ref = gr.File(label="Download") ref_btn.click(reference_align, [ref_input, mov_input, mode_dropdown_ref], [ref_image, ref_slider, reg_image, reg_slider, download_ref]) ref_slider.change(lambda i: ref_frames[i] if 0 <= i < len(ref_frames) else None, ref_slider, ref_image) reg_slider.change(lambda i: reg_frames[i] if 0 <= i < len(reg_frames) else None, reg_slider, reg_image) gr.Button("🔄 Reset Tab").click( reset_reference_based, outputs=[ref_input, mov_input, ref_image, ref_slider, reg_image, reg_slider, download_ref] ) with gr.Tab("🧩 Frame-to-Frame Alignment"): with gr.Row(): frame_file = gr.File(label="Upload Stack") ref_idx = gr.Slider(label="Reference Frame", minimum=0, maximum=0, value=0, step=1) mov_idx = gr.Slider(label="Moving Frame", minimum=0, maximum=0, value=0, step=1) gr.Examples( examples=[ ["https://github.com/glichtner/pystackreg/raw/master/examples/data/pc12-unreg.tif"] ], inputs=[frame_file], label="Try with Example Stack" ) with gr.Row(): show_adv_ftf = gr.Checkbox(label="Show Advanced Settings", value=False) mode_dropdown_ftf = gr.Dropdown(["TRANSLATION", "RIGID_BODY", "SCALED_ROTATION", "AFFINE", "BILINEAR"], value="RIGID_BODY", visible=False, label="Transformation Mode") show_adv_ftf.change(lambda v: gr.update(visible=v), show_adv_ftf, mode_dropdown_ftf) frame_btn = gr.Button("▶️ Register Frame") frame_output = gr.Image(label="Registered Output") download_ftf = gr.File(label="Download") frame_file.change( lambda f: [gr.update(value=0, maximum=len(iio.mimread(f)) - 1)] * 2 if f else [gr.update(value=0, maximum=0)] * 2, frame_file, [ref_idx, mov_idx] ) frame_btn.click(frame_to_frame_align, [frame_file, ref_idx, mov_idx, mode_dropdown_ftf], [frame_output, download_ftf]) gr.Button("🔄 Reset Tab").click( reset_frame_to_frame, outputs=[frame_file, ref_idx, mov_idx, frame_output, download_ftf] ) @demo.load( outputs=[file_input, ref_slider, frame_file, ref_idx, mov_idx, ref_input, mov_input] ) def load_from_query(request: gr.Request): params = request.query_params results = [None] * 7 # 7 outputs # One-stack file case (for ref-based + frame-to-frame) if "file_url" in params: try: url = params["file_url"] if _is_demo_url(url): tmp_path = _demo_path_for_url(url) if not os.path.exists(tmp_path): urllib.request.urlretrieve(url, tmp_path) else: tmp_path = tempfile.mktemp(suffix=".tif") # goes to WORK_DIR urllib.request.urlretrieve(url, tmp_path) stack = iio.mimread(tmp_path) max_frame = len(stack) - 1 results[0] = tmp_path # file_input results[1] = gr.update(value=0, maximum=max_frame) # ref_slider results[2] = tmp_path # frame_file results[3] = gr.update(value=0, maximum=max_frame) # ref_idx results[4] = gr.update(value=1 if max_frame >= 1 else 0, maximum=max_frame) # mov_idx except Exception as e: print(f"[Error loading file_url] {e}") # Two-stack file case (for stack-based alignment) if "file_url_1" in params and "file_url_2" in params: try: u1, u2 = params["file_url_1"], params["file_url_2"] if _is_demo_url(u1): tmp_path_1 = _demo_path_for_url(u1) if not os.path.exists(tmp_path_1): urllib.request.urlretrieve(u1, tmp_path_1) else: tmp_path_1 = tempfile.mktemp(suffix=".tif") # goes to WORK_DIR urllib.request.urlretrieve(u1, tmp_path_1) if _is_demo_url(u2): tmp_path_2 = _demo_path_for_url(u2) if not os.path.exists(tmp_path_2): urllib.request.urlretrieve(u2, tmp_path_2) else: tmp_path_2 = tempfile.mktemp(suffix=".tif") # goes to WORK_DIR urllib.request.urlretrieve(u2, tmp_path_2) results[5] = tmp_path_1 # ref_input results[6] = tmp_path_2 # mov_input except Exception as e: print(f"[Error loading file_url_1 or file_url_2] {e}") return results if __name__ == "__main__": demo.launch()