interactively delete dynamo particles - for Tom D, super hacky

tom_picker.py

from enum import Enum
from pathlib import Path
from glob import glob

import dynamotable
import mrcfile
import pandas as pd

import napari
import numpy as np
from scipy.spatial.transform import Rotation as R

from magicgui import magicgui, widgets
#/Volumes/cephfs2/tdendooven/Human_spindles/centrioles_combined/dynamo_combined/centrosome_picking/centriole_repick_test2/results/ite_0004/averages/refined_table_ref_001_ite_0004_filtered.tbl
#/Volumes/cephfs2/tdendooven/hsperm/20230421_hsperm_Krios3/movies/reconstruction/centrosome/centrosome_sperm3/results/ite_0004/averages/refined_table_ref_001_ite_0004.tbl


TABLE_FILE = '/Users/alisterburt/programming/tom-table-thing/refined_table_ref_001_ite_0004.tbl'
TABLEMAP_FILE = '/Users/alisterburt/programming/tom-table-thing/indices_column20.doc'
TOMOGRAM_PATTERN = '/Volumes/tdendooven/Flagella/movies/reconstruction/deconv/*.mrc'
LOAD_TOMOGRAMS = False

df = dynamotable.read(TABLE_FILE, TABLEMAP_FILE)

# match tomogram files to particles
if LOAD_TOMOGRAMS is True:
    deconv_tomogram_files = glob(TOMOGRAM_PATTERN)  # list of tomograms in deconv folder
    name_to_deconv_tomo_file = {Path(deconv_tomo_file).name: deconv_tomo_file for deconv_tomo_file in deconv_tomogram_files}
    tomo_name = df['tomo_file'].apply(lambda x: Path(x).name)
    deconv_tomo_file = df['tomo_name'].apply(lambda x: name_to_deconv_tomo_file[x])

df = df.drop(columns='tomo_file', errors='ignore')
# group particles
grouped = df.groupby('tomo')
valid_tomogram_files = [str(k) for k in grouped.groups.keys()]
Tomogram = Enum('Tomogram', ' '.join(valid_tomogram_files))


def get_particle_positions(df: pd.DataFrame) -> np.ndarray:
    zyx = df[['z', 'y', 'x']].to_numpy() + df[['dz', 'dy', 'dx']].to_numpy()
    return zyx


def get_particle_orientations(df: pd.DataFrame) -> np.ndarray:
    euler_angles = df[['tdrot', 'tilt', 'narot']].to_numpy()
    particle_orientations = R.from_euler('zxz', euler_angles, degrees=True).as_matrix()
    return particle_orientations[:, ::-1, ::-1]


def prepend_column(array: np.ndarray, value: float) -> np.ndarray:
    new_column = np.ones(len(array)) * value
    return np.c_[new_column, array]


viewer = napari.Viewer(ndisplay=3)


@magicgui(auto_call=True)
def add_tomogram_and_particles(tomogram: Tomogram):
    if LOAD_TOMOGRAMS is True:
        volume = mrcfile.read(tomogram.name)
    particles = grouped.get_group(int(tomogram.name))
    idx = np.arange(len(particles))
    zyx = get_particle_positions(particles)

    # orientations
    z_direction = get_particle_orientations(particles)[:, 0]
    y_direction = get_particle_orientations(particles)[:, 1]
    x_direction = get_particle_orientations(particles)[:, 2]

    z_vectors = np.stack([zyx, z_direction], axis=1)
    y_vectors = np.stack([zyx, y_direction], axis=1)
    x_vectors = np.stack([zyx, x_direction], axis=1)

    if LOAD_TOMOGRAMS is True:
        if 'tomogram' not in viewer.layers:
            viewer.add_image(data=volume, name='tomogram', colormap='gray_r')
        else:
            viewer.layers['tomogram'].data = volume

    if 'z_vectors' not in viewer.layers:
        viewer.add_vectors(z_vectors, name='z_vectors', length=10, edge_color='blue')
    else:
        viewer.layers['z_vectors'].data = z_vectors

    if 'y_vectors' not in viewer.layers:
        viewer.add_vectors(y_vectors, name='y_vectors', length=10, edge_color='yellow')
    else:
        viewer.layers['y_vectors'].data = y_vectors

    if 'x_vectors' not in viewer.layers:
        viewer.add_vectors(x_vectors, name='x_vectors', length=10, edge_color='red')
    else:
        viewer.layers['x_vectors'].data = x_vectors

    if 'particles' not in viewer.layers:
        viewer.add_points(
            zyx,
            name='particles',
            size=4,
            face_color='orange',
            features={'idx': idx},
            metadata={'particles': particles}
        )
    else:
        viewer.layers['particles'].data = zyx
        viewer.layers['particles'].features['idx'] = idx
        viewer.layers['particles'].metadata['particles'] = particles


def combine_subsetted_tables():
    output_dir = Path('output')
    output_dir.mkdir(exist_ok=True, parents=True)
    tomogram_files = {
        int(p.stem): dynamotable.read(p)
        for p in output_dir.glob('*.tbl')
    }
    tables = []
    for tomo_id, _df in grouped:
        if tomo_id not in tomogram_files:
            _df = _df.reset_index(drop=True)
            tables.append(_df)
        else:
            _df = tomogram_files[tomo_id]
            _df = _df.reset_index(drop=True)
            tables.append(_df)
    n_cols=999
    cols=[]
    for df in tables:
        if len(df.columns) < n_cols:
            n_cols=len(df.columns)
            cols=df.columns

    dfs=[_df[cols] for _df in tables]
    df_final = pd.concat(dfs)
    dynamotable.write(df_final, 'combined.tbl')





def save_particles():
    output_dir = Path('output')
    output_dir.mkdir(exist_ok=True, parents=True)
    valid_idx = np.array(viewer.layers['particles'].features['idx'])
    df = viewer.layers['particles'].metadata['particles'].iloc[valid_idx]
    df = df.drop(columns=['tomo_file'], errors='ignore')
    tomo_id = int(df['tomo'].iloc[0])
    output_file = output_dir / f'{tomo_id}.tbl'
    dynamotable.write(df, output_file)
    combine_subsetted_tables()


save_button = widgets.PushButton(text='save particles')
save_button.clicked.connect(save_particles)
add_tomogram_and_particles.append(save_button)
viewer.window.add_dock_widget(add_tomogram_and_particles)
napari.run()