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
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()