learning to project and back project a 3D object

project_backproject.py

import torch_fourier_slice as tfs
import torch
import mrcfile
import numpy as np
from scipy.spatial.transform import Rotation as R

def process_mrc(input_path, output_projection, output_reconstruction, angles=(30, 20, 18)):
    # Load and convert MRC to tensor
    volume = torch.tensor(mrcfile.read(input_path))
    
    # Create rotation matrix
    rotation = R.from_euler("ZYZ", angles, degrees=True)
    rotation_matrix = torch.tensor(rotation.as_matrix(), dtype=torch.float32)
    
    # Project 3D to 2D
    projected = tfs.project_3d_to_2d(volume, rotation_matrix)
    mrcfile.write(output_projection, np.float32(projected.cpu().detach().numpy()))
    
    # Reconstruct 3D from 2D
    rotations = rotation_matrix.T.unsqueeze(0)
    images = projected.unsqueeze(0)
    reconstruction = tfs.backproject_2d_to_3d(
        images=images.to(torch.float32),
        rotation_matrices=rotations.to(torch.float32)
    )
    
    mrcfile.write(output_reconstruction, np.float32(reconstruction.numpy()), overwrite=True)
    
    return projected, reconstruction

# Usage
projected, reconstruction = process_mrc(
    "rubsico_filtsharp.mrc",
    "projection.mrc",
    "reconstruction.mrc"
)