extract_image_subset.py

## Make a p
import argparse
from recovar import heterogeneity_volume, utils, locres
import os
import numpy as np
import logging
logger = logging.getLogger(__name__)

## To figure out which point to sample 

def center_of_mass(array):
    # Ensure the input is a NumPy array
    array = np.array(array)
    
    # Get the total mass (sum of all elements in the array)
    total_mass = np.sum(array)
    
    # Generate an array of indices for each dimension
    coords = np.indices(array.shape)
    
    # Calculate the center of mass for each dimension
    com = [np.sum(coords[i] * array) / total_mass for i in range(array.ndim)]
    
    return tuple(com)

import numpy as np

def nearest_point_index(point, points_list):
    # Convert point and points_list to NumPy arrays for easier manipulation
    point = np.array(point)
    points_array = np.array(points_list)
    
    # Calculate the Euclidean distance between the point and each point in the list
    distances = np.linalg.norm(points_array - point, axis=1)
    
    # Find the index of the nearest point
    nearest_index = np.argmin(distances)
    
    # Return the nearest point
    return nearest_index


def decide_subvolume_index_from_mask(mask, sampling_points):
    # Load the mask
    center_of_mass_mask = center_of_mass(mask)
    # Find the nearest point in the sampling points to the center of mass
    nearest_point_index = nearest_point_index(center_of_mass_mask, sampling_points)
    return nearest_point_index



def extract_image_subset(input_dir, output_path, subvolume_idx, mask, coordinate):
    if mask is not None:
        mask = utils.load_mrc(mask)
        coordinate = center_of_mass(mask)

    if coordinate is not None:
        params = utils.pickle_load(input_dir + '/params.pkl')
        # load locres?
        grid_size = utils.load_mrc(input_dir + "/locres.mrc").shape[0]
        sampling_points = locres.get_sampling_points(grid_size, params['locres_sampling'], params['locres_maskrad'], params['voxel_size']) + grid_size//2
        subvolume_idx = nearest_point_index(coordinate, sampling_points)
        logger.info(f"Extract images for feature at coordinate = {coordinate} (pixels), Nearest point to the coordinate is {subvolume_idx}")
    # heterogeneity_volume.get_inds_for_subvolume_and_save(input_dir, subvolume_idx, output_path)
    good_indices = heterogeneity_volume.get_inds_for_subvolume(input_dir, subvolume_idx)
    logger.info(f"Found {good_indices.size} images. Saving indices of subset of images")
    utils.pickle_dump(good_indices, output_path)


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Extract a subset of images from a directory')
    parser.add_argument('input_dir',type=os.path.abspath, help='A path to a folder generated by one of the script that makes volume (analyze, compute_state, compute_trajectory) (should look like /vol****/),')
    parser.add_argument( '-o', '--output', type=os.path.abspath, help='output path to a .pkl file containing the indices of subset of images')
    parser.add_argument( '-s', '--subvol-idx', dest = 'subvol_idx', type=int, help='Which subvolume is used')
    parser.add_argument(
        "--mask", metavar="mrc", default = None, help="Small mask around the feature for which you want to images. Note that this is equivalent to providing the coordinate of the center of mass of the mask." 
    )
    def list_of_ints(arg):
        return list(map(int, arg.split(',')))

    parser.add_argument( '--coordinate', default = None, dest = 'coordinate', type=list_of_ints, help="Coordinate in pixel of the feature for which you want to images. E.g. 20,30,50" )

    # parser.add_argument( '--mask', dest = 'mask', type=int, help='Which subvolume is used')
    args = parser.parse_args()
    # Check that either subvol_idx, or mask or coordinate are provided
    if args.subvol_idx is None and args.mask is None and args.coordinate is None:
        raise ValueError("You need to provide either a subvolume index, a mask or a coordinate")
    # Also check that only one is provided
    if args.subvol_idx is not None and (args.mask is not None or args.coordinate is not None):
        raise ValueError("You need to provide only one of subvolume index, mask or coordinate")

    extract_image_subset(args.input_dir, args.output, args.subvol_idx, args.mask, args.coordinate)



def test_extract_subset():




    return