starbench_processing.py

# -*- coding: utf-8 -*-
"""
High-level processing functions for STAR-Bench. (Refactored)

Handles iterating through image sets and pre-built registration pipeline configurations.
*** MODIFIED to accept and pass through pipeline_map ***
*** MODIFIED to pass landmarks to generate_visualizations ***
"""

import logging
from pathlib import Path
import cv2
import numpy as np
import time
from tqdm import tqdm
# Use relative imports for the package structure
import starbench_registration as registration
import starbench_utils as utils
import starbench_reporting as reporting
import starbench_visualization as visualization
# Import classes needed for type hinting and execution
try:
    from starbench_detectors import DetectorFactory
except ImportError:
    logging.error("Unable to import DetectorFactory in starbench_processing for type hinting.")
    DetectorFactory = None


# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def _format_time(seconds):
    """Format seconds into a readable string (hours, minutes, seconds)."""
    seconds = int(seconds)
    hours, remainder = divmod(seconds, 3600)
    minutes, seconds = divmod(remainder, 60)
    if hours > 0:
        return f"{hours}h {minutes}m {seconds}s"
    elif minutes > 0:
        return f"{minutes}m {seconds}s"
    else:
        return f"{seconds}s"


def process_from_folder(folder_path: str,
                        pipeline_map: dict[str, str],
                        output_dir: str,
                        detector_factory: DetectorFactory,
                        tag: str = None,
                        visualize: bool = False,
                        device: str = 'cpu',
                        ransac_thresh: float = 5.0):
    """Process image sets from a folder, evaluating each backend individually."""
    matching_sets = utils.find_matching_files_in_folder(folder_path, tag)
    if not matching_sets:
        logger.warning(f"No matching file set found in '{folder_path}'{f' for tag {tag}' if tag else ''}.")
        return {}

    if not pipeline_map:
        logger.warning("Empty pipeline_map. No pipelines to run.")
        return {}

    total_sets = len(matching_sets)
    results_by_set = {}
    overall_start_time = time.time()

    for i, file_set in enumerate(tqdm(matching_sets, desc="Progress")):
        set_name = f"{file_set['tag']}{file_set['number']}"
        logger.info(f"\n{'='*10} Starting processing of Set {i + 1}/{total_sets}: {set_name} {'='*10}")

        set_start_time = time.time()
        set_output_dir = Path(output_dir) / set_name
        set_output_dir.mkdir(parents=True, exist_ok=True)

        try:
            # landmarks_fix are GT points on the fixed image
            # landmarks_mov are GT points on the moving image
            _, _, landmarks_fix, landmarks_mov, transform_gt = utils.load_ground_truth(file_set['gt_file'])
            if landmarks_fix is None or landmarks_mov is None or transform_gt is None:
                raise ValueError("Ground truth landmarks or transformation matrix could not be loaded.")
        except Exception as e:
            logger.error(f"Unable to load or validate ground truth for set {set_name}. Skipping set. Error: {e}", exc_info=True)
            continue

        results_by_method = {}
        logger.info(f"Ground truth loaded with {len(landmarks_fix)} fixed landmarks and {len(landmarks_mov)} moving landmarks.")

        # Iterate over the detector map (one entry per method).
        for pipeline_name in tqdm(pipeline_map.keys(), desc="Method", leave=False):
            pretty_name = pipeline_map.get(pipeline_name, pipeline_name)

            logger.info(f"\n--- Running Method: {pretty_name} (ID: {pipeline_name}) ---")

            method_results = registration.process_image_pair_with_starbench(
                moving_img_path=file_set['moving_file'],
                fixed_img_path=file_set['fixed_file'],
                gt_img_path=file_set.get('gt_image'),
                landmarks_mov=landmarks_mov,
                transform_gt=transform_gt,
                pipeline_name=pipeline_name,
                detector_factory=detector_factory,
                device=device,
                ransac_thresh=ransac_thresh
            )

            if method_results:
                # The KEY of the results is the programmatic name (must be unique)
                results_by_method[pipeline_name] = method_results

                logger.debug(f"Method '{pretty_name}' results: {method_results}")
                
                # Pass the "pretty" name for saving the metrics file
                reporting.save_metrics(method_results, pretty_name, str(set_output_dir))

                if visualize and method_results.get('final_transform') is not None:
                    # Pass the "pretty" name for the visualization filenames
                    # Pass GT landmarks for the new visualization
                    generate_visualizations(
                        file_set,
                        method_results,
                        pretty_name,
                        str(set_output_dir),
                        landmarks_mov=landmarks_mov,
                        landmarks_gt_on_fixed=landmarks_fix
                    )
            else:
                logger.error(f"Processing with {pretty_name} (ID: {pipeline_name}) failed for set {set_name}.")
                results_by_method[pipeline_name] = {
                    'failure_stage': -1,
                    'pipeline': pipeline_name,
                    'error': 'Processing function returned None'
                }


        if results_by_method:
            results_by_set[set_name] = results_by_method
            # Pass the translation dictionary to the report function
            reporting.compare_methods(results_by_method, str(set_output_dir), pipeline_map)

        set_end_time = time.time()
        set_duration = set_end_time - set_start_time
        logger.info(f"--- Set {set_name} completed in {_format_time(set_duration)} ---")

        if i == 0 and total_sets > 1:
            remaining_sets = total_sets - 1
            estimated_remaining_time = set_duration * remaining_sets
            formatted_etc = _format_time(estimated_remaining_time)
            logger.info(f"*** COMPLETION ESTIMATE: {remaining_sets} sets remaining. Estimated remaining time: {formatted_etc} ***")

    if results_by_set:
        logger.debug(f"All sets processed. Compiling final summary report. Total sets: {len(results_by_set)}")
        all_methods_run = sorted(list(set(m for res in results_by_set.values() for m in res.keys())))
        # Pass the translation dictionary to the final summary report
        logger.debug(f"Methods run across all sets: {all_methods_run}")
        reporting.create_summary_report(results_by_set, all_methods_run, output_dir, pipeline_map)
    else:
        logger.warning("No results were generated for any set.")

    overall_duration = time.time() - overall_start_time
    logger.info(f"\n{'='*10} Processing completed for all {total_sets} sets in {_format_time(overall_duration)} {'='*10}")
    
    return results_by_set


def generate_visualizations(file_set, method_results, method_name, output_dir, landmarks_mov=None, landmarks_gt_on_fixed=None):
    """Helper that emits visualisations for a single method run.

    Takes landmarks as well so the error-visualisation plot can be rendered.
    ``method_name`` is the pretty (human-readable) label.
    """
    logger.info(f"Generating visualisations for {method_name.upper()}...")
    fixed_img = cv2.imread(file_set['fixed_file'], cv2.IMREAD_GRAYSCALE)
    moving_img = cv2.imread(file_set['moving_file'], cv2.IMREAD_GRAYSCALE)

    if fixed_img is None or moving_img is None:
        logger.error(f"Unable to load images for visualisation of {method_name}")
        return

    # Attach landmarks to the result dict so the visualiser can use them.
    vis_results = {
        'transformation_matrix': method_results.get('final_transform'),
        'keypoints_fixed': method_results.get('keypoints_fixed'),
        'keypoints_moving': method_results.get('keypoints_moving'),
        'landmarks_mov': landmarks_mov,
        'landmarks_gt_on_fixed': landmarks_gt_on_fixed,
    }


    visualization.visualize_results(
        fixed_img, moving_img, vis_results,
        method_name, output_dir
    )