Stub scan image files

stub_scan_image.py

#!/usr/bin/env python3
"""
stub_scan_image_file.py
──────────────────────────────────────────────────────────────────────────────
Functions for creating stub ScanImage TIFF files with modified metadata.
"""

import re
import shutil
import struct
from pathlib import Path
import tifffile


def create_stub_scan_image(file_path, output_file_path=None, n_frames=120, new_volumes=90):
    """
    Create a stub ScanImage TIFF file with modified metadata.
    
    Parameters
    ----------
    file_path : str or Path
        Path to the original ScanImage TIFF file
    output_file_path : str or Path, optional
        Path to save the stub file. If None, will use original filename with "_stub" suffix
    n_frames : int, optional
        Number of frames to keep in the stub file
    new_volumes : int, optional
        Value to set for both actualNumVolumes and numVolumes
        
    Returns
    -------
    dict
        Dictionary containing information about the stub file:
        - 'path': Path to the stub file
        - 'original_volumes': Original volume counts
        - 'new_volumes': New volume counts
        - 'n_frames': Number of frames kept
    """
    file_path = Path(file_path)
    if not file_path.exists():
        raise FileNotFoundError(f"File not found: {file_path}")
    
    # Create stub file path if not provided
    if output_file_path is None:
        output_file_path = file_path.with_name(file_path.stem + "_stub" + file_path.suffix)
    else:
        output_file_path = Path(output_file_path)
    
    # Copy the file
    shutil.copyfile(file_path, output_file_path)
    
    # Regex for finding volume counts in metadata
    vol_regex = re.compile(
        rb"(SI\.hStackManager\.(?:actualNumVolumes|numVolumes)\s*=\s*)(\d+)",
        re.I,
    )
    
    def replace_volumes(block):
        """Replace volume counts in block with new_volumes."""
        for match in vol_regex.finditer(block):
            start, end = match.span(2)
            width = end - start
            block[start:end] = f"{new_volumes:>{width}}".encode()
        return block
    
    # Get original metadata and IFD info
    with tifffile.TiffFile(output_file_path) as tiff:
        first_ifd = tiff.pages[0].offset
        
        # Get original volume values
        original_volumes = {}
        if 'SI.hStackManager.actualNumVolumes' in tiff.scanimage_metadata['FrameData']:
            original_volumes['actualNumVolumes'] = tiff.scanimage_metadata['FrameData']['SI.hStackManager.actualNumVolumes']
        if 'SI.hStackManager.numVolumes' in tiff.scanimage_metadata['FrameData']:
            original_volumes['numVolumes'] = tiff.scanimage_metadata['FrameData']['SI.hStackManager.numVolumes']
    
    # Modify the stub file
    with output_file_path.open("r+b") as f:
        # Read and modify metadata
        f.seek(0)
        tiff_header = f.read(16)
        
        # Read ScanImage metadata section
        magic_number = struct.unpack('<I', f.read(4))[0]
        si_version = struct.unpack('<I', f.read(4))[0]
        non_varying_length = struct.unpack('<I', f.read(4))[0]
        roi_group_length = struct.unpack('<I', f.read(4))[0]
        
        # Modify non-varying metadata
        non_varying_data_offset = f.tell()
        non_varying_data = bytearray(f.read(non_varying_length))
        non_varying_data = replace_volumes(non_varying_data)
        f.seek(non_varying_data_offset)
        f.write(non_varying_data)
        
        # Break IFD chain after n_frames and get position for truncation
        offset = first_ifd
        next_frame_pos = None
        for i in range(n_frames):
            f.seek(offset)
            n_entries = struct.unpack("<Q", f.read(8))[0]
            next_ptr_pos = offset + 8 + n_entries * 20
            f.seek(next_ptr_pos)
            next_ifd = struct.unpack("<Q", f.read(8))[0]
            if i == n_frames - 1:
                f.seek(next_ptr_pos)
                f.write(struct.pack("<Q", 0))  # terminate chain
                next_frame_pos = next_ifd  # Save position for truncation
                break
            offset = next_ifd
        
        # Truncate the file to remove excess frames if we found the position
        if next_frame_pos:
            f.seek(next_frame_pos)
            f.truncate()
    
    # Verify pixel data
    with tifffile.TiffFile(file_path) as orig, tifffile.TiffFile(output_file_path) as stub:
        for idx in range(n_frames):
            if not (orig.pages[idx].asarray() == stub.pages[idx].asarray()).all():
                raise ValueError(f"Pixel data mismatch at frame {idx}")
    
    # Verify metadata changes
    with tifffile.TiffFile(output_file_path) as tiff:
        metadata = tiff.scanimage_metadata['FrameData']
        new_values = {}
        if 'SI.hStackManager.actualNumVolumes' in metadata:
            new_values['actualNumVolumes'] = metadata['SI.hStackManager.actualNumVolumes']
        if 'SI.hStackManager.numVolumes' in metadata:
            new_values['numVolumes'] = metadata['SI.hStackManager.numVolumes']
    
    # Calculate file sizes in MiB
    original_size_mib = file_path.stat().st_size / (1024 * 1024)
    stub_size_mib = output_file_path.stat().st_size / (1024 * 1024)
    size_reduction_percent = (1 - (stub_size_mib / original_size_mib)) * 100
    
    # Return information about the stub file
    return {
        'path': output_file_path,
        'original_volumes': original_volumes,
        'new_volumes': new_values,
        'n_frames': n_frames,
        'original_size_mib': original_size_mib,
        'stub_size_mib': stub_size_mib,
        'size_reduction_percent': size_reduction_percent
    }


# Example usage
if __name__ == "__main__":
    # Default settings
    file_path = Path.home() / "Downloads" / "vol_no_flyback_00001_00001.tif"
    n_frames = 90
    new_volumes = 10
    
    try:
        result = create_stub_scan_image(file_path, n_frames=n_frames, new_volumes=new_volumes)
        
        print(f"Created stub file: {result['path']}")
        print(f"Original volumes: {result['original_volumes']}")
        print(f"New volumes: {result['new_volumes']}")
        print(f"Kept {result['n_frames']} frames")
        print(f"Original file size: {result['original_size_mib']:.2f} MiB")
        print(f"Stub file size: {result['stub_size_mib']:.2f} MiB")
        print(f"Size reduction: {result['size_reduction_percent']:.2f}%")
        
    except Exception as e:
        print(f"Error: {e}")