app.py

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

from core.utils import *

# 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"]
            tmp_path = tempfile.mktemp(suffix=".tif")
            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"]
                tmp_path = tempfile.mktemp(suffix=".tif")
                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:
                tmp_path_1 = tempfile.mktemp(suffix=".tif")
                tmp_path_2 = tempfile.mktemp(suffix=".tif")
                urllib.request.urlretrieve(params["file_url_1"], tmp_path_1)
                urllib.request.urlretrieve(params["file_url_2"], 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()