align.py

import numpy as np
from PIL import Image
import os
import matplotlib
matplotlib.use('Agg')  # Set this before importing pyplot
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # This is needed for 3D plotting
from concurrent.futures import ThreadPoolExecutor
import multiprocessing
import copy
import json


def calculate_scale_factor(png_dimensions, obj_bounds):
    """
    Calculate scaling factor between PNG and OBJ spaces
    
    Args:
        png_dimensions: (width, height) of PNG in pixels
        obj_bounds: ((min_x, max_x), (min_y, max_y), (min_z, max_z)) of OBJ model
    """
    # Get physical dimensions of OBJ model
    obj_width = obj_bounds[0][1] - obj_bounds[0][0]
    obj_height = obj_bounds[1][1] - obj_bounds[1][0]
    
    # Calculate scaling factors for both dimensions
    scale_x = obj_width / png_dimensions[0]
    scale_y = obj_height / png_dimensions[1]
    
    # Use average or minimum scale to maintain aspect ratio
    return min(scale_x, scale_y)

def convert_mm_to_obj_space(slice_position_mm, obj_bounds, physical_bounds_mm):
    """
    Convert millimeter position to OBJ coordinate space
    
    Args:
        slice_position_mm: Position of slice in mm
        obj_bounds: Bounds of OBJ model
        physical_bounds_mm: Physical dimensions in mm
    """
    # Calculate conversion factor from mm to OBJ units
    obj_height = obj_bounds[2][1] - obj_bounds[2][0]
    physical_height = physical_bounds_mm[2][1] - physical_bounds_mm[2][0]
    mm_to_obj = obj_height / physical_height
    
    # Convert position
    return obj_bounds[2][0] + (slice_position_mm * mm_to_obj)

def register_slice(png_slice, z_coordinate, scale_factor, obj_bounds):
    """
    Register PNG slice at the correct position in 3D space with proper centering
    """
    slice_array = np.array(png_slice)
    height, width = slice_array.shape[:2]
    
    # Create coordinate grids
    y, x = np.mgrid[:height, :width]
    
    # Scale coordinates
    x = x * scale_factor
    y = y * scale_factor
    
    # Center the slice relative to OBJ model
    x = x + obj_bounds[0][0]  # Offset by min X of OBJ
    y = y + obj_bounds[1][0]  # Offset by min Y of OBJ
    z = np.full_like(x, z_coordinate)
    
    return np.stack([x, y, z], axis=-1)

def inspect_and_calibrate_obj(obj_model):
    """
    Inspect OBJ dimensions and apply Allen reference space offset
    Returns a new offset copy that matches reference space
    """
    scaled_model = copy.deepcopy(obj_model)
    vertices = np.array(scaled_model['vertices'])
    
    # Calculate current bounds
    min_bounds = vertices.min(axis=0)
    max_bounds = vertices.max(axis=0)
    
    # Just store vertices back without scaling
    scaled_model['vertices'] = vertices
    
    # Calculate bounds
    physical_bounds_mm = (
        (vertices[:,0].min(), vertices[:,0].max()),
        (vertices[:,1].min(), vertices[:,1].max()),
        (vertices[:,2].min(), vertices[:,2].max())
    )
    
    return scaled_model, physical_bounds_mm

def visualize_alignment(obj_model, aligned_coords, png_slice):
    """
    Enhanced visualization to better show slice position in 3D space
    """
    try:
        # Create 3D plot
        fig = plt.figure(figsize=(15, 10))
        ax = fig.add_subplot(111, projection='3d')
        
        # Plot OBJ model vertices (subsample for visualization)
        vertices = np.array(obj_model['vertices'])
        step = max(1, len(vertices) // 1000)  # Subsample for better visualization
        ax.scatter(vertices[::step,0], vertices[::step,1], vertices[::step,2], 
                  c='gray', alpha=0.5, s=5, label='OBJ Model')  # Increased point size
        
        # Plot PNG slice as points
        ax.scatter(aligned_coords[:,0], aligned_coords[:,1], aligned_coords[:,2],
                  c='red', alpha=0.8, s=10, label='PNG Slice')  # Increased point size and opacity
        
        # Set axis limits explicitly
        ax.set_xlim([vertices[:,0].min(), vertices[:,0].max()])
        ax.set_ylim([vertices[:,1].min(), vertices[:,1].max()])
        ax.set_zlim([vertices[:,2].min(), vertices[:,2].max()])
        
        # Add labels and legend
        ax.set_xlabel('X (mm)')
        ax.set_ylabel('Y (mm)')
        ax.set_zlabel('Z (mm)')
        ax.legend()
        
        # Adjust view angle for better visibility
        ax.view_init(elev=30, azim=45)
        
        # Save visualization with higher DPI
        plt.savefig('alignment_visualization.png', dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"Saved visualization to alignment_visualization.png")
        
        # Print statistics
        print("\nAlignment Statistics:")
        print(f"OBJ Model bounds (mm):")
        print(f"X: [{vertices[:,0].min():.2f}, {vertices[:,0].max():.2f}]")
        print(f"Y: [{vertices[:,1].min():.2f}, {vertices[:,1].max():.2f}]")
        print(f"Z: [{vertices[:,2].min():.2f}, {vertices[:,2].max():.2f}]")
        
        print(f"\nPNG Slice bounds (mm):")
        print(f"X: [{aligned_coords[:,0].min():.2f}, {aligned_coords[:,0].max():.2f}]")
        print(f"Y: [{aligned_coords[:,1].min():.2f}, {aligned_coords[:,1].max():.2f}]")
        print(f"Z: [{aligned_coords[:,2].min():.2f}, {aligned_coords[:,2].max():.2f}]")
        
        # Print point counts for debugging
        print(f"\nPoint counts:")
        print(f"OBJ vertices: {len(vertices)}")
        print(f"PNG points: {len(aligned_coords)}")
        
    except Exception as e:
        print(f"Error in visualization: {e}")
        import traceback
        traceback.print_exc()

def load_sample_info(reference_dir):
    """Load and parse sampleInfo.json"""
    info_path = os.path.join(reference_dir, 'sampleInfo.json')
    with open(info_path, 'r') as f:
        return json.load(f)

def calculate_slice_positions(obj_model, sample_info, orientation):
    """
    Calculate slice positions using sampleInfo parameters
    
    Args:
        obj_model: Dictionary containing OBJ model data
        sample_info: Parsed sampleInfo.json data
        orientation: 'horizontal', 'sagittal', or 'coronal'
    """
    vertices = np.array(obj_model['vertices'])
    pixel_size = np.array(sample_info['pixelSize'])
    origin = np.array(sample_info['origin'])
    display_range = sample_info['displayRange']
    
    # Normalize coordinates to match Allen reference space
    # First subtract origin, then divide by pixel size
    vertices = (vertices + origin) / pixel_size
    
    # Get bounds in slice space
    min_bounds = vertices.min(axis=0)
    max_bounds = vertices.max(axis=0)
    
    if orientation == 'horizontal':
        slice_coord = 2  # Z determines slice
        in_slice_coords = [0, 1]  # X,Y within slice
    elif orientation == 'sagittal':
        slice_coord = 0  # X determines slice
        in_slice_coords = [1, 2]  # Y,Z within slice
    else:  # coronal
        slice_coord = 1  # Y determines slice
        in_slice_coords = [0, 2]  # X,Z within slice
    
    # Map to slice indices within display range
    slice_min = max(display_range[0], int(min_bounds[slice_coord]))
    slice_max = min(display_range[1], int(max_bounds[slice_coord]))
    
    print(f"\n{orientation.capitalize()} slices:")
    print(f"Range: {slice_min} to {slice_max}")
    print(f"Slice coordinate ({['X','Y','Z'][slice_coord]}): {min_bounds[slice_coord]:.2f} to {max_bounds[slice_coord]:.2f}")
    print(f"In-slice coordinates ({['X','Y','Z'][in_slice_coords[0]]},{['X','Y','Z'][in_slice_coords[1]]}): "
          f"({min_bounds[in_slice_coords[0]]:.2f},{min_bounds[in_slice_coords[1]]:.2f}) to "
          f"({max_bounds[in_slice_coords[0]]:.2f},{max_bounds[in_slice_coords[1]]:.2f})")
    
    return slice_min, slice_max, min_bounds[in_slice_coords], max_bounds[in_slice_coords]

def align_png_to_obj(png_slice, obj_model, slice_position_mm, orientation='horizontal', visualize=True):
    """
    Align PNG slice with OBJ model using Allen reference space offset
    """
    # Load sample info for offset
    sample_info = load_sample_info('sample/allen-reference')
    origin = np.array(sample_info['origin'])
    
    # Apply offset when converting coordinates
    vertices = np.array(obj_model['vertices'])
    vertices = vertices + origin  # Add offset to match reference space
    
    # Calibrate OBJ model to millimeters
    obj_model, physical_bounds_mm = inspect_and_calibrate_obj(obj_model)
    
    # Get model bounds
    obj_bounds = (
        (vertices[:,0].min(), vertices[:,0].max()),
        (vertices[:,1].min(), vertices[:,1].max()),
        (vertices[:,2].min(), vertices[:,2].max())
    )
    
    # Calculate scale factor
    png_dimensions = png_slice.size
    scale_factor = calculate_scale_factor(png_dimensions, obj_bounds)
    
    # Convert slice position to OBJ space
    z_coordinate = convert_mm_to_obj_space(slice_position_mm, obj_bounds, physical_bounds_mm)
    
    # Register slice
    aligned_coords = register_slice(png_slice, z_coordinate, scale_factor, obj_bounds)
    
    # Add visualization
    if visualize:
        visualize_alignment(obj_model, aligned_coords, png_slice)
    
    return aligned_coords

def load_obj_file(file_path):
    """
    Load an OBJ file and return a dictionary containing the model data
    
    Args:
        file_path: Path to the OBJ file
        
    Returns:
        Dictionary containing:
            - vertices: list of (x,y,z) coordinates
            - faces: list of vertex indices for each face
    """
    vertices = []
    faces = []
    
    try:
        with open(file_path, 'r') as f:
            for line in f:
                if line.startswith('#'): # Skip comments
                    continue
                    
                values = line.split()
                if not values:
                    continue
                
                if values[0] == 'v':  # Vertex
                    # Convert vertex coordinates to float
                    vertex = [float(values[1]), float(values[2]), float(values[3])]
                    vertices.append(vertex)
                    
                elif values[0] == 'f':  # Face
                    # Handle different face formats (v, v/vt, v/vt/vn)
                    face = []
                    for v in values[1:]:
                        # Get the vertex index (before any '/' character)
                        vertex_idx = int(v.split('/')[0]) - 1  # OBJ indices start at 1
                        face.append(vertex_idx)
                    faces.append(face)
    
    except Exception as e:
        print(f"Error loading OBJ file: {e}")
        return None
    
    return {
        'vertices': vertices,
        'faces': faces
    }

def process_all_slices(obj_model, input_dir, output_dir):
    """Process all slices with memory optimization"""
    def process_slice(orientation, png_file, scale_factor):
        input_path = os.path.join(input_dir, orientation.capitalize(), png_file)
        output_path = os.path.join(output_dir, orientation.capitalize(), png_file)
        
        if os.path.exists(input_path):
            try:
                png_slice = Image.open(input_path)
                new_size = tuple(int(dim * scale_factor) for dim in png_slice.size)
                # Process in smaller chunks if image is too large
                if max(new_size) > 10000:  # Threshold for large images
                    print(f"Warning: Large scaling for {png_file}, may need adjustment")
                    scale_factor = scale_factor * 0.1  # Reduce scaling
                    new_size = tuple(int(dim * scale_factor) for dim in png_slice.size)
                
                png_slice = png_slice.resize(new_size, Image.LANCZOS)
                os.makedirs(os.path.dirname(output_path), exist_ok=True)
                png_slice.save(output_path)
                png_slice.close()  # Explicitly close to free memory
                return f"Processed {png_file} for {orientation}"
            except Exception as e:
                return f"Error processing {png_file}: {str(e)}"
    
    # Calculate scale factors once
    scale_factors = {}
    for orientation in ['horizontal', 'sagittal', 'coronal']:
        _, _, scale_factor = calculate_slice_positions(obj_model, orientation)
        scale_factors[orientation] = scale_factor
    
    # Process each orientation separately to manage memory
    for orientation in scale_factors:
        print(f"\nProcessing {orientation} slices...")
        input_orientation_dir = os.path.join(input_dir, orientation.capitalize())
        if os.path.exists(input_orientation_dir):
            png_files = sorted([f for f in os.listdir(input_orientation_dir) if f.endswith('.png')])
            
            # Process files in smaller batches
            batch_size = 10
            for i in range(0, len(png_files), batch_size):
                batch = png_files[i:i + batch_size]
                print(f"Processing batch {i//batch_size + 1}/{len(png_files)//batch_size + 1}")
                
                for png_file in batch:
                    result = process_slice(orientation, png_file, scale_factors[orientation])
                    print(result)
    
    # Copy additional files
    for file in ['Allen.nrrd', 'sampleInfo.json']:
        input_path = os.path.join(input_dir, file)
        output_path = os.path.join(output_dir, file)
        if os.path.exists(input_path):
            import shutil
            shutil.copy2(input_path, output_path)
            print(f"Copied {file} to {output_path}")

def save_obj_file(obj_model, file_path):
    """
    Save an OBJ model to a file
    
    Args:
        obj_model: Dictionary containing vertices and faces
        file_path: Path to save the OBJ file
    """
    try:
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        with open(file_path, 'w') as f:
            # Write vertices
            for v in obj_model['vertices']:
                f.write(f"v {v[0]:.6f} {v[1]:.6f} {v[2]:.6f}\n")
            
            # Write faces
            for face in obj_model['faces']:
                # OBJ indices start at 1, not 0
                face_str = ' '.join(str(idx + 1) for idx in face)
                f.write(f"f {face_str}\n")
                
        print(f"Saved scaled OBJ to {file_path}")
        
    except Exception as e:
        print(f"Error saving OBJ file: {e}")

if __name__ == "__main__":
    # Load OBJ and reference info
    obj_model = load_obj_file('test_obj/obj/997.obj')
    sample_info = load_sample_info('sample/allen-reference')
    
    # Apply Allen reference space offset
    vertices = np.array(obj_model['vertices'])
    origin = np.array(sample_info['origin'])
    obj_model['vertices'] = vertices + origin
    
    # Save the offset model
    save_obj_file(obj_model, 'test_obj/obj/997_aligned.obj')