<commit_before>#!/usr/bin/env python


import sys
import os
import os.path
import numpy
import shutil
import logging
import json
from skimage import morphology as skmorph
from scipy.ndimage import label
import traceback
import glob
import re
import h5py

from . import imio, agglo, morpho, classify, app_logger, \
    session_manager, pixel, features

try:
    from gala import stack_np
except ImportError:
    np_installed = False   
else:
    np_installed = True

try:
    import syngeo
except ImportError:
    logging.warning('Could not import syngeo. ' +
                                        'Synapse-aware mode not available.')

# Group where we store predictions in HDF5 file
PREDICTIONS_HDF5_GROUP = '/volume/predictions'


def grab_boundary(prediction, channels, master_logger):
    boundary = None
    master_logger.debug("Grabbing boundary labels: " + str(channels))
    for channel_id in channels:
        if boundary is None:
            boundary = prediction[...,channel_id] 
        else:
            boundary += prediction[...,channel_id]

    return boundary


def gen_supervoxels(options, prediction_file, master_logger):
    """Returns ndarray labeled using (optionally seeded) watershed algorithm

    Args:
        options:  OptionNamespace.
        prediction_file:  String.  File name of prediction hdf5 file where predictions
            are assumed to be in group PREDICTIONS_HDF5_GROUP.

    Returns:
        A 2-tuple of supervoxel and prediction ndarray.
    """
    master_logger.debug("Generating supervoxels")
    if not os.path.isfile(prediction_file):
        raise Exception("Training file not found: " + prediction_file)

    prediction = imio.read_image_stack(prediction_file, group=PREDICTIONS_HDF5_GROUP)
    master_logger.info("Transposed boundary prediction")
    prediction = prediction.transpose((2, 1, 0, 3))

    #if options.extract_ilp_prediction:
    #   prediction = prediction.transpose((2, 1, 0))

    # TODO -- Refactor.  If 'single-channel' and hdf5 prediction file is given, it looks like
    #   read_image_stack will return a modified volume and the bound-channels parameter must
    #   be 0 or there'll be conflict.
    boundary = grab_boundary(prediction, options.bound_channels, master_logger) 
    master_logger.info("Shape of boundary: %s" % str(boundary.shape))

    # Prediction file is in format (t, x, y, z, c) but needs to be in format (z, x, y).
    # Also, raveler convention is (0,0) sits in bottom left while ilastik convention is
    # origin sits in top left.
    # imio.read_image_stack squeezes out the first dim.

    master_logger.debug("watershed seed value threshold: " + str(options.seed_val))
    seeds = label(boundary<=options.seed_val)[0]

    if options.seed_size > 0:
        master_logger.debug("Removing small seeds")
        seeds = morpho.remove_small_connected_components(seeds, options.seed_size)
        master_logger.debug("Finished removing small seeds")

    master_logger.info("Starting watershed")
    
    boundary_cropped = boundary
    seeds_cropped = seeds 
    if options.border_size > 0:
        boundary_cropped = boundary[options.border_size:(-1*options.border_size), options.border_size:(-1*options.border_size),options.border_size:(-1*options.border_size)]
        seeds_cropped = seeds[options.border_size:(-1*options.border_size), options.border_size:(-1*options.border_size),options.border_size:(-1*options.border_size)]

    # Returns a matrix labeled using seeded watershed
    supervoxels_cropped = skmorph.watershed(boundary_cropped, seeds_cropped)
    
    supervoxels = supervoxels_cropped
    if options.border_size > 0:
        supervoxels = seeds.copy()
        supervoxels.dtype = supervoxels_cropped.dtype
        supervoxels[:,:,:] = 0 
        supervoxels[options.border_size:(-1*options.border_size), options.border_size:(-1*options.border_size),options.border_size:(-1*options.border_size)] = supervoxels_cropped

    master_logger.info("Finished watershed")

    if options.synapse_file is not None:
        master_logger.info("Processing synapses")
        pre_post_pairs = syngeo.io.raveler_synapse_annotations_to_coords(
            options.synapse_file)
        synapse_volume = syngeo.io.volume_synapse_view(pre_post_pairs, boundary.shape)
        if options.border_size > 0:
            synvol_cropped = synapse_volume[options.border_size:(-1*options.border_size), options.border_size:(-1*options.border_size),options.border_size:(-1*options.border_size)]
            synvol_cropped = synvol_cropped.copy()
            synapse_volume[:,:,:] = 0
            synapse_volume[options.border_size:(-1*options.border_size), options.border_size:(-1*options.border_size),options.border_size:(-1*options.border_size)] = synvol_cropped
             
        
        supervoxels = morpho.split_exclusions(boundary, supervoxels, synapse_volume, 
                                                        options.synapse_dilation)
        master_logger.info("Finished processing synapses")
   
    return supervoxels, prediction


def agglomeration(options, agglom_stack, supervoxels, prediction, 
        image_stack, session_location, sp_outs, master_logger):
    
    seg_thresholds = sorted(options.segmentation_thresholds)
    for threshold in seg_thresholds:
        if threshold != 0 or not options.use_neuroproof:
            master_logger.info("Starting agglomeration to threshold " + str(threshold)
                + " with " + str(agglom_stack.number_of_nodes()))
            agglom_stack.agglomerate(threshold)
            master_logger.info("Finished agglomeration to threshold " + str(threshold)
                + " with " + str(agglom_stack.number_of_nodes()))
            
            if options.inclusion_removal:
                inclusion_removal(agglom_stack, master_logger)

        segmentation = agglom_stack.get_segmentation()     

        if options.h5_output:
            imio.write_image_stack(segmentation,
                session_location+"/agglom-"+str(threshold)+".lzf.h5", compression='lzf')
          
        
        file_base = os.path.abspath(session_location)+"/seg_data/agglom-"+str(threshold)
        transforms = imio.compute_sp_to_body_map(supervoxels, segmentation)
        seg_loc = file_base +"v1.h5"
        if not os.path.exists(session_location+"/seg_data"):
            os.makedirs(session_location+"/seg_data")
        imio.write_mapped_segmentation(supervoxels, transforms, seg_loc)    

        if options.synapse_file is not None:
            h5temp = h5py.File(seg_loc, 'a')
            syn_data = json.load(open((options.synapse_file)))
            meta = syn_data['metadata']
            meta['username'] = "auto"
            syn_data_str = json.dumps(syn_data, indent=4)
            h5temp.create_dataset("synapse-annotations", data=syn_data_str)

        graph_loc = file_base+"-graphv1.json"
        agglom_stack.write_plaza_json(graph_loc, options.synapse_file)
       
        json_data = {}
        json_data['graph'] = graph_loc
        json_data['border'] = options.border_size  
        subvolume = {}
        subvolume['segmentation-file'] = seg_loc
        
        gray_file_whole = os.path.abspath(glob.glob(options.image_stack)[0])
        gray_path = os.path.dirname(gray_file_whole)
       
        gray_file = os.path.basename(gray_file_whole)
        field_width = len(re.findall(r'\d',gray_file))
        field_rep = "%%0%dd" % field_width
        gray_file = re.sub(r'\d+', field_rep, gray_file)
        
        subvolume['grayscale-files'] = gray_path + "/" + gray_file
        
        # get extant
        x1 = 0
        y1 = 0
        z1 = 0
        z2,y2,x2 = supervoxels.shape
        z2 = z2 - 2*options.border_size
        y2 = y2 - 2*options.border_size
        x2 = x2 - 2*options.border_size
        extant = re.findall(r'\d+-\d+_\d+-\d+_\d+-\d+', gray_path)
        if len(extant) > 0:
            bbox = extant[0]
            x1,x2,y1,y2,z1,z2 = re.findall(r'\d+', bbox)
        subvolume["far-upper-right"] = [x2,y2,z2]
        subvolume["near-lower-left"] = [x1,y1,z1]

        json_data['subvolumes'] = [subvolume]
         
        # write out json file 
        json_file = session_location + "/segmentation-" + str(threshold) + "-v1.json"
        jw = open(json_file, 'w')
        jw.write(json.dumps(json_data, indent=4))

        #if options.raveler_output:
        #    sps_outs = output_raveler(segmentation, supervoxels, image_stack, "agglom-" + str(threshold),
        #        session_location, master_logger)   
        #    master_logger.info("Writing graph.json")
        #    agglom_stack.write_plaza_json(session_location+"/raveler-export/agglom-"+str(threshold)+"/graph.json",
        #                                    options.synapse_file)
        #    if options.synapse_file is not None:
        #        shutil.copyfile(options.synapse_file,
        #                session_location + "/raveler-export/agglom-"+str(threshold)+"/annotations-synapse.json") 
        #    master_logger.info("Finished writing graph.json")


def inclusion_removal(agglom_stack, master_logger):
    master_logger.info("Starting inclusion removal with " + str(agglom_stack.number_of_nodes()) + " nodes")
    agglom_stack.remove_inclusions()
    master_logger.info("Finished inclusion removal with " + str(agglom_stack.number_of_nodes()) + " nodes")


def output_raveler(segmentation, supervoxels, grayscale, name, session_location, master_logger, sps_out=None):
    """Output segmented data to a Raveler formatted directory

    Args:
        segmentation:  ndarray.<commit_msg>updates output format<commit_after>#!/usr/bin/env python


import sys
import os
import os.path
import numpy
import shutil
import logging
import json
from skimage import morphology as skmorph
from scipy.ndimage import label
import traceback
import glob
import re
import h5py

from . import imio, agglo, morpho, classify, app_logger, \
    session_manager, pixel, features

try:
    from gala import stack_np
except ImportError:
    np_installed = False   
else:
    np_installed = True

try:
    import syngeo
except ImportError:
    logging.warning('Could not import syngeo. ' +
                                        'Synapse-aware mode not available.')

# Group where we store predictions in HDF5 file
PREDICTIONS_HDF5_GROUP = '/volume/predictions'


def grab_boundary(prediction, channels, master_logger):
    boundary = None
    master_logger.debug("Grabbing boundary labels: " + str(channels))
    for channel_id in channels:
        if boundary is None:
            boundary = prediction[...,channel_id] 
        else:
            boundary += prediction[...,channel_id]

    return boundary


def gen_supervoxels(options, prediction_file, master_logger):
    """Returns ndarray labeled using (optionally seeded) watershed algorithm

    Args:
        options:  OptionNamespace.
        prediction_file:  String.  File name of prediction hdf5 file where predictions
            are assumed to be in group PREDICTIONS_HDF5_GROUP.

    Returns:
        A 2-tuple of supervoxel and prediction ndarray.
    """
    master_logger.debug("Generating supervoxels")
    if not os.path.isfile(prediction_file):
        raise Exception("Training file not found: " + prediction_file)

    prediction = imio.read_image_stack(prediction_file, group=PREDICTIONS_HDF5_GROUP)
    master_logger.info("Transposed boundary prediction")
    prediction = prediction.transpose((2, 1, 0, 3))

    #if options.extract_ilp_prediction:
    #   prediction = prediction.transpose((2, 1, 0))

    # TODO -- Refactor.  If 'single-channel' and hdf5 prediction file is given, it looks like
    #   read_image_stack will return a modified volume and the bound-channels parameter must
    #   be 0 or there'll be conflict.
    boundary = grab_boundary(prediction, options.bound_channels, master_logger) 
    master_logger.info("Shape of boundary: %s" % str(boundary.shape))

    # Prediction file is in format (t, x, y, z, c) but needs to be in format (z, x, y).
    # Also, raveler convention is (0,0) sits in bottom left while ilastik convention is
    # origin sits in top left.
    # imio.read_image_stack squeezes out the first dim.

    master_logger.debug("watershed seed value threshold: " + str(options.seed_val))
    seeds = label(boundary<=options.seed_val)[0]

    if options.seed_size > 0:
        master_logger.debug("Removing small seeds")
        seeds = morpho.remove_small_connected_components(seeds, options.seed_size)
        master_logger.debug("Finished removing small seeds")

    master_logger.info("Starting watershed")
    
    boundary_cropped = boundary
    seeds_cropped = seeds 
    if options.border_size > 0:
        boundary_cropped = boundary[options.border_size:(-1*options.border_size), options.border_size:(-1*options.border_size),options.border_size:(-1*options.border_size)]
        seeds_cropped = seeds[options.border_size:(-1*options.border_size), options.border_size:(-1*options.border_size),options.border_size:(-1*options.border_size)]

    # Returns a matrix labeled using seeded watershed
    supervoxels_cropped = skmorph.watershed(boundary_cropped, seeds_cropped)
    
    supervoxels = supervoxels_cropped
    if options.border_size > 0:
        supervoxels = seeds.copy()
        supervoxels.dtype = supervoxels_cropped.dtype
        supervoxels[:,:,:] = 0 
        supervoxels[options.border_size:(-1*options.border_size), options.border_size:(-1*options.border_size),options.border_size:(-1*options.border_size)] = supervoxels_cropped

    master_logger.info("Finished watershed")

    if options.synapse_file is not None:
        master_logger.info("Processing synapses")
        pre_post_pairs = syngeo.io.raveler_synapse_annotations_to_coords(
            options.synapse_file)
        synapse_volume = syngeo.io.volume_synapse_view(pre_post_pairs, boundary.shape)
        if options.border_size > 0:
            synvol_cropped = synapse_volume[options.border_size:(-1*options.border_size), options.border_size:(-1*options.border_size),options.border_size:(-1*options.border_size)]
            synvol_cropped = synvol_cropped.copy()
            synapse_volume[:,:,:] = 0
            synapse_volume[options.border_size:(-1*options.border_size), options.border_size:(-1*options.border_size),options.border_size:(-1*options.border_size)] = synvol_cropped
             
        
        supervoxels = morpho.split_exclusions(boundary, supervoxels, synapse_volume, 
                                                        options.synapse_dilation)
        master_logger.info("Finished processing synapses")
   
    return supervoxels, prediction


def agglomeration(options, agglom_stack, supervoxels, prediction, 
        image_stack, session_location, sp_outs, master_logger):
    
    seg_thresholds = sorted(options.segmentation_thresholds)
    for threshold in seg_thresholds:
        if threshold != 0 or not options.use_neuroproof:
            master_logger.info("Starting agglomeration to threshold " + str(threshold)
                + " with " + str(agglom_stack.number_of_nodes()))
            agglom_stack.agglomerate(threshold)
            master_logger.info("Finished agglomeration to threshold " + str(threshold)
                + " with " + str(agglom_stack.number_of_nodes()))
            
            if options.inclusion_removal:
                inclusion_removal(agglom_stack, master_logger)

        segmentation = agglom_stack.get_segmentation()     

        if options.h5_output:
            imio.write_image_stack(segmentation,
                session_location+"/agglom-"+str(threshold)+".lzf.h5", compression='lzf')
          
        
        file_base = os.path.abspath(session_location)+"/seg_data/agglom-"+str(threshold)
        transforms = imio.compute_sp_to_body_map(supervoxels, segmentation)
        seg_loc = file_base +"v1.h5"
        if not os.path.exists(session_location+"/seg_data"):
            os.makedirs(session_location+"/seg_data")
        imio.write_mapped_segmentation(supervoxels, transforms, seg_loc)    

        if options.synapse_file is not None:
            h5temp = h5py.File(seg_loc, 'a')
            syn_data = json.load(open((options.synapse_file)))
            meta = syn_data['metadata']
            meta['username'] = "auto"
            syn_data_str = json.dumps(syn_data, indent=4)
            h5temp.create_dataset("synapse-annotations", data=syn_data_str)

        graph_loc = file_base+"-graphv1.json"
        agglom_stack.write_plaza_json(graph_loc, options.synapse_file)
       
        json_data = {}
        json_data['graph'] = graph_loc
        json_data['border'] = options.border_size  
        subvolume = {}
        subvolume['segmentation-file'] = seg_loc
        
        gray_file_whole = os.path.abspath(glob.glob(options.image_stack)[0])
        gray_path = os.path.dirname(gray_file_whole)
       
        gray_file = os.path.basename(gray_file_whole)
        field_width = len(re.findall(r'\d',gray_file))
        field_rep = "%%0%dd" % field_width
        gray_file = re.sub(r'\d+', field_rep, gray_file)
        
        subvolume['grayscale-files'] = gray_path + "/" + gray_file
        
        # get extant
        x1 = 0
        y1 = 0
        z1 = 0
        z2,y2,x2 = supervoxels.shape
        z2 = z2 - 2*options.border_size
        y2 = y2 - 2*options.border_size
        x2 = x2 - 2*options.border_size
        extant = re.findall(r'\d+-\d+_\d+-\d+_\d+-\d+', gray_path)
        if len(extant) > 0:
            bbox = extant[0]
            x1,x2,y1,y2,z1,z2 = re.findall(r'\d+', bbox)
        subvolume["far-upper-right"] = [int(x2),int(y2),int(z2)]
        subvolume["near-lower-left"] = [int(x1),int(y1),int(z1)]

        json_data['subvolumes'] = [subvolume]
         
        # write out json file 
        json_file = session_location + "/segmentation-" + str(threshold) + "-v1.json"
        jw = open(json_file, 'w')
        jw.write(json.dumps(json_data, indent=4))

        #if options.raveler_output:
        #    sps_outs = output_raveler(segmentation, supervoxels, image_stack, "agglom-" + str(threshold),
        #        session_location, master_logger)   
        #    master_logger.info("Writing graph.json")
        #    agglom_stack.write_plaza_json(session_location+"/raveler-export/agglom-"+str(threshold)+"/graph.json",
        #                                    options.synapse_file)
        #    if options.synapse_file is not None:
        #        shutil.copyfile(options.synapse_file,
        #                session_location + "/raveler-export/agglom-"+str(threshold)+"/annotations-synapse.json") 
        #    master_logger.info("Finished writing graph.json")


def inclusion_removal(agglom_stack, master_logger):
    master_logger.info("Starting inclusion removal with " + str(agglom_stack.number_of_nodes()) + " nodes")
    agglom_stack.remove_inclusions()
    master_logger.info("Finished inclusion removal with " + str(agglom_stack.number_of_nodes()) + " nodes")


def output_raveler(segmentation, supervoxels, grayscale, name, session_location, master_logger, sps_out=None):
    """Output segmented data to a Raveler formatted directory

    Args:
        segmentation:  ndarray.
{'fraction_non_alphanumeric': 0.10228928199791884, 'fraction_numerical': 0.00894901144640999, 'mean_word_length': 2.937922556853104, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4712919', 'n_tokens_mistral': 6125, 'n_tokens_neox': 5844, 'n_words': 1147}