lattice-sub 1.0.10__py3-none-any.whl

lattice_subtraction/processing.py

@@ -0,0 +1,221 @@
+"""
+Image processing utilities.
+
+This module contains functions for image padding, background subtraction,
+and other preprocessing operations.
+"""
+
+import numpy as np
+from typing import Tuple
+
+
+def pad_image(
+    image: np.ndarray,
+    pad_origin: Tuple[int, int],
+    target_size: int | None = None,
+    pad_value: float | None = None,
+) -> Tuple[np.ndarray, dict]:
+    """
+    Pad an image with mean border for FFT processing.
+    
+    This replicates the MATLAB padarray functionality with 'pre' and 'post'
+    padding using the image mean value.
+    
+    Args:
+        image: Input 2D image
+        pad_origin: Padding offsets (pad_y, pad_x) - pixels to add at start
+        target_size: Target square size. If None, auto-calculated.
+        pad_value: Value to use for padding. If None, uses image mean.
+        
+    Returns:
+        Tuple of:
+        - Padded image
+        - Metadata dict with original shape and padding info for cropping
+    """
+    orig_h, orig_w = image.shape
+    pad_y, pad_x = pad_origin
+    
+    # Auto-calculate target size if not provided
+    if target_size is None:
+        max_dim = max(orig_h, orig_w)
+        target_size = max_dim + pad_x * 2
+        # Round to nearest 10 for FFT efficiency
+        target_size = int(np.round(target_size / 10) * 10)
+    
+    # Calculate padding amounts
+    pad_top = pad_y - 1 if pad_y > 0 else 0
+    pad_left = pad_x - 1 if pad_x > 0 else 0
+    pad_bottom = target_size - orig_h - pad_top
+    pad_right = target_size - orig_w - pad_left
+    
+    # Ensure non-negative padding
+    pad_bottom = max(0, pad_bottom)
+    pad_right = max(0, pad_right)
+    
+    # Use image mean for padding value if not specified
+    if pad_value is None:
+        pad_value = float(np.mean(image))
+    
+    # Perform padding
+    padded = np.pad(
+        image,
+        ((pad_top, pad_bottom), (pad_left, pad_right)),
+        mode='constant',
+        constant_values=pad_value,
+    )
+    
+    # Store metadata for later cropping
+    metadata = {
+        'original_shape': (orig_h, orig_w),
+        'pad_top': pad_top,
+        'pad_left': pad_left,
+        'pad_bottom': pad_bottom,
+        'pad_right': pad_right,
+        'target_size': target_size,
+    }
+    
+    return padded, metadata
+
+
+def crop_to_original(
+    image: np.ndarray,
+    metadata: dict,
+) -> np.ndarray:
+    """
+    Crop a padded image back to its original size.
+    
+    Args:
+        image: Padded image
+        metadata: Metadata dict from pad_image()
+        
+    Returns:
+        Cropped image with original dimensions
+    """
+    orig_h, orig_w = metadata['original_shape']
+    pad_top = metadata['pad_top']
+    pad_left = metadata['pad_left']
+    
+    return image[pad_top:pad_top + orig_h, pad_left:pad_left + orig_w]
+
+
+def subtract_background(
+    image: np.ndarray,
+    median_filter_size: int = 10,
+) -> np.ndarray:
+    """
+    Subtract smooth background from an image to reveal sharp features.
+    
+    This is the Python equivalent of bg_FastSubtract_standard.m.
+    It creates a smoothed version of the image using median filtering
+    and subtracts it from the original.
+    
+    Args:
+        image: Input 2D image (typically log-power spectrum)
+        median_filter_size: Size of median filter kernel. Default: 10
+        
+    Returns:
+        Background-subtracted image with edge regions replaced by mean
+    """
+    from scipy import ndimage
+    from skimage.transform import resize
+    
+    h, w = image.shape
+    
+    if max(h, w) < 500:
+        # For small images, apply median filter directly
+        smoothed = ndimage.median_filter(image, size=median_filter_size)
+        edge = median_filter_size
+    else:
+        # For large images, downsample -> filter -> upsample
+        shrink_factor = 500 / max(h, w)
+        
+        # Downsample
+        small = resize(
+            image, 
+            (int(h * shrink_factor), int(w * shrink_factor)),
+            order=1,  # Bilinear
+            preserve_range=True,
+            anti_aliasing=True,
+        )
+        
+        # Apply median filter
+        small = ndimage.median_filter(small, size=median_filter_size)
+        
+        # Upsample back to original size
+        smoothed = resize(
+            small,
+            (h, w),
+            order=1,
+            preserve_range=True,
+            anti_aliasing=True,
+        )
+        
+        # Scale edge hiding region
+        edge = int(median_filter_size / shrink_factor)
+    
+    # Subtract background
+    subtracted = image - smoothed
+    
+    # Hide edges with artifacts
+    mean_value = np.mean(subtracted)
+    edge = max(1, edge)
+    
+    # Replace edge regions with mean
+    subtracted[:edge, :] = mean_value
+    subtracted[-edge:, :] = mean_value
+    subtracted[:, :edge] = mean_value
+    subtracted[:, -edge:] = mean_value
+    
+    return subtracted.astype(np.float32)
+
+
+def compute_power_spectrum(
+    fft_shifted: np.ndarray,
+    log_scale: bool = True,
+    epsilon: float = 1e-10,
+) -> np.ndarray:
+    """
+    Compute power spectrum from shifted FFT.
+    
+    Args:
+        fft_shifted: Centered FFT (after fftshift)
+        log_scale: If True, return log of amplitude. Default: True
+        epsilon: Small value to avoid log(0). Default: 1e-10
+        
+    Returns:
+        Power spectrum (log-amplitude if log_scale=True)
+    """
+    amplitude = np.abs(fft_shifted)
+    
+    if log_scale:
+        return np.log(amplitude + epsilon)
+    
+    return amplitude
+
+
+def shift_and_average(
+    array: np.ndarray,
+    shift_pixels: int,
+) -> np.ndarray:
+    """
+    Create a local average by averaging 4 shifted copies.
+    
+    This is the inpainting technique from bg_push_by_rot.m that
+    averages amplitude values from neighboring regions.
+    
+    Args:
+        array: Input 2D array (typically FFT amplitude)
+        shift_pixels: Number of pixels to shift in each direction
+        
+    Returns:
+        Averaged array (local background estimate)
+    """
+    # Shift in 4 cardinal directions and average
+    shifted_sum = (
+        np.roll(array, shift_pixels, axis=0) +
+        np.roll(array, -shift_pixels, axis=0) +
+        np.roll(array, shift_pixels, axis=1) +
+        np.roll(array, -shift_pixels, axis=1)
+    )
+    
+    return shifted_sum / 4.0

lattice_subtraction/ui.py

@@ -0,0 +1,256 @@
+"""
+Terminal UI utilities for lattice subtraction.
+
+This module provides styled terminal output with ASCII art banner
+and formatted progress messages. Output is only shown when running
+interactively (TTY detected) and not suppressed by --quiet flag.
+
+When piped or used in a pipeline, decorative output is automatically
+suppressed to avoid polluting downstream processing.
+"""
+
+import sys
+import time
+from typing import Optional
+
+
+class Colors:
+    """ANSI color codes for terminal styling."""
+    HEADER = '\033[95m'
+    BLUE = '\033[94m'
+    CYAN = '\033[96m'
+    GREEN = '\033[92m'
+    YELLOW = '\033[93m'
+    RED = '\033[91m'
+    BOLD = '\033[1m'
+    DIM = '\033[2m'
+    RESET = '\033[0m'
+
+
+# ASCII Art Banner
+BANNER = r"""
+.__          __    __  .__                                   ___.    
+|  | _____ _/  |__/  |_|__| ____  ____             ________ _\_ |__  
+|  | \__  \\   __\   __\  |/ ___\/ __ \   ______  /  ___/  |  \ __ \ 
+|  |__/ __ \|  |  |  | |  \  \__\  ___/  /_____/  \___ \|  |  / \_\ \
+|____(____  /__|  |__| |__|\___  >___  >         /____  >____/|___  /
+          \/                   \/    \/               \/          \/ 
+"""
+
+# Import version from package to keep it in sync
+from . import __version__ as VERSION
+
+
+def is_interactive() -> bool:
+    """
+    Check if running in an interactive terminal.
+    
+    Returns False if stdout is piped or redirected, which means
+    we're likely part of a pipeline and should suppress decorative output.
+    """
+    return sys.stdout.isatty()
+
+
+class TerminalUI:
+    """
+    Manages styled terminal output for the CLI.
+    
+    Decorative output (banner, colors, progress indicators) is only shown when:
+    - Running in an interactive terminal (TTY detected)
+    - Not suppressed by quiet mode
+    
+    When piped or in a script, output is automatically minimal.
+    """
+    
+    def __init__(self, quiet: bool = False):
+        """
+        Initialize the terminal UI.
+        
+        Args:
+            quiet: If True, suppress all decorative output even in interactive mode
+        """
+        self.quiet = quiet
+        self.interactive = is_interactive() and not quiet
+        self.use_colors = self.interactive
+        self._start_time: Optional[float] = None
+        self._file_start_time: Optional[float] = None
+    
+    def _colorize(self, text: str, color: str) -> str:
+        """Apply color if colors are enabled."""
+        if self.use_colors:
+            return f"{color}{text}{Colors.RESET}"
+        return text
+    
+    def print_banner(self) -> None:
+        """Print the ASCII art banner."""
+        if not self.interactive:
+            return
+        
+        print()
+        print(self._colorize(BANNER, Colors.CYAN))
+        tagline = f"  Phase-preserving FFT inpainting for cryo-EM  |  v{VERSION}"
+        print(self._colorize(tagline, Colors.DIM))
+        print()
+    
+    def print_config(self, pixel_size: float, threshold: float, 
+                     backend: str, gpu_name: Optional[str] = None) -> None:
+        """Print configuration summary."""
+        if not self.interactive:
+            return
+        
+        print(self._colorize("  Configuration", Colors.BOLD))
+        print(self._colorize("  -------------", Colors.DIM))
+        print(f"    Pixel size:  {self._colorize(f'{pixel_size} A', Colors.YELLOW)}")
+        print(f"    Threshold:   {self._colorize(str(threshold), Colors.YELLOW)}")
+        
+        # Determine backend display string
+        if backend == "auto":
+            # Check if GPU is actually available for auto mode
+            if gpu_name:
+                backend_str = f"Auto → GPU ({gpu_name})"
+                print(f"    Backend:     {self._colorize(backend_str, Colors.GREEN)}")
+            else:
+                backend_str = "Auto → CPU"
+                print(f"    Backend:     {self._colorize(backend_str, Colors.BLUE)}")
+        elif backend == "pytorch" and gpu_name:
+            backend_str = f"PyTorch CUDA ({gpu_name})"
+            print(f"    Backend:     {self._colorize(backend_str, Colors.GREEN)}")
+        elif backend == "pytorch":
+            print(f"    Backend:     {self._colorize('PyTorch CUDA', Colors.GREEN)}")
+        else:
+            print(f"    Backend:     {self._colorize('NumPy (CPU)', Colors.BLUE)}")
+        print()
+    
+    def start_timer(self) -> None:
+        """Start the overall timer."""
+        self._start_time = time.time()
+    
+    def start_processing(self, filename: str, shape: Optional[tuple] = None) -> None:
+        """Indicate start of processing a file."""
+        self._file_start_time = time.time()
+        
+        if not self.interactive:
+            return
+        
+        print(f"  {self._colorize('>', Colors.CYAN)} Processing: {self._colorize(filename, Colors.BOLD)}")
+        if shape:
+            print(f"    {self._colorize('|-', Colors.DIM)} Size: {shape[0]} x {shape[1]}")
+    
+    def end_processing(self, output_path: str, success: bool = True) -> None:
+        """Indicate end of processing."""
+        elapsed = time.time() - self._file_start_time if self._file_start_time else 0
+        
+        if not self.interactive:
+            return
+        
+        if success:
+            status = self._colorize(f"[OK] Complete ({elapsed:.2f}s)", Colors.GREEN)
+        else:
+            status = self._colorize("[FAIL]", Colors.RED)
+        
+        print(f"    {self._colorize('`-', Colors.DIM)} {status}")
+        print()
+    
+    def print_batch_header(self, num_files: int, output_dir: str, num_workers: int = 1) -> None:
+        """Print batch processing header."""
+        if not self.interactive:
+            return
+        
+        print(self._colorize("  Batch Processing", Colors.BOLD))
+        print(self._colorize("  ----------------", Colors.DIM))
+        print(f"    Files:    {self._colorize(str(num_files), Colors.YELLOW)}")
+        print(f"    Output:   {output_dir}")
+        print(f"    Workers:  {num_workers}")
+        print()
+    
+    def print_batch_progress(self, current: int, total: int, filename: str, 
+                             elapsed: Optional[float] = None) -> None:
+        """Print batch progress update."""
+        if not self.interactive:
+            return
+        
+        progress = f"[{current}/{total}]"
+        time_str = f" ({elapsed:.1f}s)" if elapsed else ""
+        print(f"  {self._colorize(progress, Colors.CYAN)} {filename}{time_str}")
+    
+    def print_batch_complete(self) -> None:
+        """Print batch completion message."""
+        if not self.interactive:
+            return
+        
+        elapsed = time.time() - self._start_time if self._start_time else 0
+        print()
+        print(f"  {self._colorize('[OK]', Colors.GREEN)} {self._colorize('Batch complete', Colors.BOLD)} ({elapsed:.1f}s)")
+        print()
+    
+    def print_summary(self, processed: int, failed: int = 0) -> None:
+        """Print final summary."""
+        if not self.interactive:
+            return
+        
+        elapsed = time.time() - self._start_time if self._start_time else 0
+        
+        print(self._colorize("  Summary", Colors.BOLD))
+        print(self._colorize("  -------", Colors.DIM))
+        print(f"    Processed: {self._colorize(str(processed), Colors.GREEN)}")
+        if failed > 0:
+            print(f"    Failed:    {self._colorize(str(failed), Colors.RED)}")
+        print(f"    Time:      {elapsed:.1f}s")
+        print()
+    
+    def print_error(self, message: str) -> None:
+        """Print error message (always shown unless quiet)."""
+        if self.quiet:
+            return
+        prefix = self._colorize("Error:", Colors.RED) if self.use_colors else "Error:"
+        print(f"  {prefix} {message}", file=sys.stderr)
+    
+    def print_warning(self, message: str) -> None:
+        """Print warning message."""
+        if not self.interactive:
+            return
+        prefix = self._colorize("Warning:", Colors.YELLOW)
+        print(f"  {prefix} {message}", file=sys.stderr)
+    
+    def print_success(self, message: str) -> None:
+        """Print success message."""
+        if not self.interactive:
+            return
+        check = self._colorize("[OK]", Colors.GREEN)
+        print(f"  {check} {message}")
+    
+    def print_info(self, message: str) -> None:
+        """Print info message."""
+        if not self.interactive:
+            return
+        print(f"  {self._colorize('>', Colors.CYAN)} {message}")
+    
+    def print_saved(self, path: str) -> None:
+        """Print saved file notification."""
+        if not self.interactive:
+            return
+        print(f"    {self._colorize('|-', Colors.DIM)} Saved: {path}")
+
+
+def get_ui(quiet: bool = False) -> TerminalUI:
+    """
+    Get a TerminalUI instance.
+    
+    Args:
+        quiet: If True, suppress decorative output even in interactive mode
+        
+    Returns:
+        TerminalUI instance configured for current environment
+    """
+    return TerminalUI(quiet=quiet)
+
+
+def get_gpu_name() -> Optional[str]:
+    """Get the name of the CUDA GPU if available."""
+    try:
+        import torch
+        if torch.cuda.is_available():
+            return torch.cuda.get_device_name(0)
+    except ImportError:
+        pass
+    return None

lattice_subtraction/visualization.py

@@ -0,0 +1,195 @@
+"""
+Visualization utilities for lattice subtraction results.
+
+This module provides functions to create comparison visualizations
+showing original, processed, and difference images.
+"""
+
+import logging
+from pathlib import Path
+from typing import Optional, Tuple
+
+import numpy as np
+
+# Silence matplotlib's verbose debug logging
+logging.getLogger('matplotlib').setLevel(logging.WARNING)
+logging.getLogger('matplotlib.font_manager').setLevel(logging.WARNING)
+
+logger = logging.getLogger(__name__)
+
+
+def create_comparison_figure(
+    original: np.ndarray,
+    processed: np.ndarray,
+    title: str = "Lattice Subtraction Comparison",
+    figsize: Tuple[int, int] = (18, 6),
+    dpi: int = 150,
+):
+    """
+    Create a comparison figure showing original, processed, and difference images.
+    
+    Args:
+        original: Original image array
+        processed: Processed (lattice-subtracted) image array
+        title: Figure title
+        figsize: Figure size in inches (width, height)
+        dpi: Resolution for saving
+        
+    Returns:
+        matplotlib Figure object
+    """
+    import matplotlib.pyplot as plt
+    
+    # Compute difference
+    difference = original - processed
+    
+    # Create figure
+    fig, axes = plt.subplots(1, 3, figsize=figsize)
+    
+    # Contrast limits from original
+    vmin, vmax = np.percentile(original, [1, 99])
+    
+    # Original
+    axes[0].imshow(original, cmap='gray', vmin=vmin, vmax=vmax)
+    axes[0].set_title(f'Original\n{original.shape}')
+    axes[0].axis('off')
+    
+    # Lattice Subtracted
+    axes[1].imshow(processed, cmap='gray', vmin=vmin, vmax=vmax)
+    axes[1].set_title(f'Lattice Subtracted\n{processed.shape}')
+    axes[1].axis('off')
+    
+    # Difference (removed lattice)
+    diff_std = np.std(difference)
+    axes[2].imshow(
+        difference, 
+        cmap='RdBu_r',
+        vmin=-diff_std * 3, 
+        vmax=diff_std * 3
+    )
+    axes[2].set_title('Difference (Removed Lattice)')
+    axes[2].axis('off')
+    
+    # Title
+    plt.suptitle(title, fontsize=14)
+    plt.tight_layout()
+    
+    return fig
+
+
+def save_comparison_visualization(
+    original_path: Path,
+    processed_path: Path,
+    output_path: Path,
+    dpi: int = 150,
+) -> None:
+    """
+    Create and save a comparison visualization for a single image pair.
+    
+    Args:
+        original_path: Path to original MRC file
+        processed_path: Path to processed MRC file
+        output_path: Path for output PNG file
+    """
+    import matplotlib.pyplot as plt
+    import mrcfile
+    
+    # Load images
+    with mrcfile.open(original_path, 'r') as f:
+        original = f.data.copy()
+    with mrcfile.open(processed_path, 'r') as f:
+        processed = f.data.copy()
+    
+    # Create title
+    name = original_path.name
+    short_name = name[:60] + "..." if len(name) > 60 else name
+    title = f"Lattice Subtraction Comparison: {short_name}"
+    
+    # Create and save figure
+    fig = create_comparison_figure(original, processed, title=title, dpi=dpi)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(output_path, dpi=dpi, bbox_inches='tight')
+    plt.close(fig)
+
+
+def generate_visualizations(
+    input_dir: Path,
+    output_dir: Path,
+    viz_dir: Path,
+    prefix: str = "sub_",
+    pattern: str = "*.mrc",
+    dpi: int = 150,
+    show_progress: bool = True,
+) -> Tuple[int, int]:
+    """
+    Generate comparison visualizations for all processed images in a directory.
+    
+    Args:
+        input_dir: Directory containing original MRC files
+        output_dir: Directory containing processed MRC files
+        viz_dir: Directory for output visualization PNG files
+        prefix: Prefix used for processed files (default: "sub_")
+        pattern: Glob pattern for finding processed files
+        dpi: Resolution for output images
+        show_progress: Show progress bar
+        
+    Returns:
+        Tuple of (successful_count, total_count)
+    """
+    import matplotlib.pyplot as plt
+    
+    viz_dir = Path(viz_dir)
+    viz_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Find all processed files
+    output_files = sorted(Path(output_dir).glob(f"{prefix}{pattern}"))
+    
+    if not output_files:
+        logger.warning(f"No processed files found matching '{prefix}{pattern}' in {output_dir}")
+        return 0, 0
+    
+    successful = 0
+    total = len(output_files)
+    
+    # Setup iterator with optional progress bar
+    if show_progress:
+        try:
+            from tqdm import tqdm
+            iterator = tqdm(output_files, desc="Generating visualizations", unit="file")
+        except ImportError:
+            iterator = output_files
+    else:
+        iterator = output_files
+    
+    for processed_path in iterator:
+        try:
+            # Get corresponding input file
+            input_name = processed_path.name.replace(prefix, "", 1)
+            input_path = Path(input_dir) / input_name
+            
+            if not input_path.exists():
+                logger.debug(f"Original not found: {input_path}")
+                continue
+            
+            # Output path
+            viz_name = input_name.replace(".mrc", ".png")
+            viz_path = viz_dir / viz_name
+            
+            # Skip if already exists
+            if viz_path.exists():
+                successful += 1
+                continue
+            
+            # Generate visualization
+            save_comparison_visualization(
+                original_path=input_path,
+                processed_path=processed_path,
+                output_path=viz_path,
+                dpi=dpi,
+            )
+            successful += 1
+            
+        except Exception as e:
+            logger.error(f"Failed to create visualization for {processed_path.name}: {e}")
+    
+    return successful, total