deepivus/cli.py

"""CLI entrypoints for DeepIVUS."""

import datetime
import os
from typing import Optional

import click
from deepivus.config import resolve_bifurcation_threshold


@click.group()
def cli() -> None:
    """DeepIVUS CLI."""


@cli.command("segment")
@click.argument("dicom_path", type=click.Path(exists=True, dir_okay=False))
@click.option(
    "--output-prefix",
    "-o",
    default=None,
    type=str,
    help="Output path prefix (without extension). Defaults to output/<timestamp>/<input filename>.",
)
@click.option(
    "--fps",
    default=None,
    type=float,
    help="Overlay video FPS. Defaults to DICOM CineRate or 30.",
)
@click.option(
    "--bifurcation-threshold",
    default=None,
    show_default=False,
    type=click.FloatRange(min=0.0, max=1.0),
    help="Threshold for bifurcation classifier labels. Defaults to threshold.json beside the selected bifurcation model.",
)
@click.option(
    "--framewise/--no-framewise",
    default=False,
    show_default=True,
    help="Run inference frame-by-frame (batch_size=1) to simulate realtime processing.",
)
def segment_cmd(
    dicom_path: str,
    output_prefix: Optional[str],
    fps: Optional[float],
    bifurcation_threshold: Optional[float],
    framewise: bool,
) -> None:
    """Segment lumen and classify bifurcation for all frames."""
    from .pipeline import segment_and_export

    if output_prefix is None:
        stem = os.path.splitext(os.path.basename(dicom_path))[0]
        output_root = os.path.join(os.getcwd(), "output")
        os.makedirs(output_root, exist_ok=True)

        run_folder = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        output_dir = os.path.join(output_root, run_folder)
        suffix = 1
        while os.path.exists(output_dir):
            output_dir = os.path.join(output_root, f"{run_folder}_{suffix}")
            suffix += 1

        os.makedirs(output_dir, exist_ok=True)
        output_prefix = os.path.join(output_dir, stem)
    else:
        output_dir = os.path.dirname(os.path.abspath(output_prefix))
        os.makedirs(output_dir, exist_ok=True)

    if bifurcation_threshold is None:
        bifurcation_threshold = resolve_bifurcation_threshold(default=0.5)

    (
        xml_path,
        json_path,
        top_conf_json_path,
        video_path,
        bif_overlay_video_path,
        bif_json_path,
        bif_summary_path,
    ) = segment_and_export(
        dicom_path,
        output_prefix,
        fps,
        bifurcation_threshold,
        framewise=framewise,
    )

    click.echo(f"Contours XML: {xml_path}")
    click.echo(f"Contours JSONL: {json_path}")
    click.echo(f"Top confidence JSONL: {top_conf_json_path}")
    click.echo(f"Overlay video: {video_path}")
    click.echo(f"Overlay video (with bifurcation flags): {bif_overlay_video_path}")
    click.echo(f"Bifurcation predictions JSONL: {bif_json_path}")
    click.echo(f"Bifurcation summary JSON: {bif_summary_path}")


@cli.command("edit-annotations")
@click.argument("dicom_path", type=click.Path(exists=True, dir_okay=False))
@click.option(
    "--annotations-path",
    default=None,
    type=click.Path(exists=True, dir_okay=False),
    help=(
        "Base contour JSONL to edit. "
        "Defaults to latest output/<timestamp>/<dicom_stem>_contours.jsonl."
    ),
)
@click.option(
    "--edits-path",
    default=None,
    type=click.Path(dir_okay=False),
    help="Path for edited annotations JSONL. Defaults beside base file as *_edited_annotations.jsonl.",
)
@click.option(
    "--output-root",
    default="output",
    show_default=True,
    type=click.Path(file_okay=False),
    help="Root folder containing pipeline outputs used for default annotation discovery.",
)
def edit_annotations_cmd(
    dicom_path: str,
    annotations_path: Optional[str],
    edits_path: Optional[str],
    output_root: str,
) -> None:
    """Open GUI editor to review and adjust contour coordinates frame-by-frame."""
    from .gui import launch_annotation_editor

    launch_annotation_editor(
        dicom_path=dicom_path,
        annotations_path=annotations_path,
        edits_path=edits_path,
        output_root=output_root,
    )