---

Segmentation.h

Go to the documentation of this file.

/***************************************************************************
 *cr
 *cr            (C) Copyright 1995-2019 The Board of Trustees of the
 *cr                        University of Illinois
 *cr                         All Rights Reserved
 *cr
 ***************************************************************************/

/***************************************************************************
 * RCS INFORMATION:
 *
 *      $RCSfile: Segmentation.h,v $
 *      $Author: johns $        $Locker:  $             $State: Exp $
 *      $Revision: 1.6 $        $Date: 2019/01/17 21:21:01 $
 *
 ***************************************************************************
 * DESCRIPTION:
 *   Scale-space variant of Watershed image segmentation intended for use
 *   on 3-D cryo-EM density maps.
 ***************************************************************************/
#ifndef SEGMENTATION_H
#define SEGMENTATION_H

#include "molfile_plugin.h"
#include "GaussianBlur.h"
#include "ScaleSpaceFilter.h"
#include "VolumetricData.h"

#define TIMER


class Segmentation {

  public:
    Segmentation(const VolumetricData *map, bool cuda=true);

    Segmentation(float* data,
                 molfile_volumetric_t* metatdata,
                 bool cuda=true);

    Segmentation(unsigned short* data,
                 molfile_volumetric_t* metatdata,
                 bool cuda=true);

    Segmentation(unsigned char* data,
                 molfile_volumetric_t* metatdata,
                 bool cuda=true);

    ~Segmentation();

    double segment(int   num_final_groups,
                   float watershed_blur_sigma,
                   float blur_initial_sigma,
                   float blur_multiple,
                   MERGE_POLICY policy,
                   const bool verbose=true);

    template <typename OUT_T>
    void get_results(OUT_T* results);

    unsigned long get_num_groups();

  private:

    // The object keeps track of the current smallest datatype
    // we can use based on the number of groups in the segmentation.
    enum Datatype {
      DT_LONG,
      DT_INT,
      DT_SHORT,
      DT_CHAR,
      DT_ULONG,
      DT_UINT,
      DT_USHORT,
      DT_UCHAR,
      DT_FLOAT,
      DT_DOUBLE,
      DT_ERROR
    };

    unsigned long height;
    unsigned long width;
    unsigned long depth;
    unsigned long nVoxels;
    unsigned long nGroups;
    Datatype group_type;
    Datatype image_type;
    GaussianBlur<float>* gaussian_f;
    GaussianBlur<unsigned short>* gaussian_us;
    GaussianBlur<unsigned char>* gaussian_uc;
    unsigned long* segments_ul;
    unsigned int* segments_ui;
    unsigned short* segments_us;
    bool verbose;


    // The object will always try to use cuda unless explicitly told not to.
    // If an error occurs at any point in the segmentation process we should be 
    // able to fall back to the CPU seamlessly (TODO fallback not implemented yet)
    bool use_cuda;

    void init(unsigned long w,
              unsigned long h,
              unsigned long d,
              bool cuda);

    void update_type();

    void switch_to_cpu();

    Datatype get_new_group_type();

    template <typename IN_T, typename OUT_T>
    void copy_and_convert_type(IN_T* in, OUT_T* out, unsigned long num_elems,
                               bool copy_from_gpu=false);

    template <typename GROUP_T>
    void get_watershed_segmentation(float watershed_blur_sigma,
                                    GROUP_T* segments);

    template <typename GROUP_T>
    void merge_groups(unsigned long num_final_groups,
                      float blur_initial_sigma,
                      float blur_multiple, 
                      MERGE_POLICY policy,
                      GROUP_T* segments);

    template <typename GROUP_T, typename IMAGE_T>
    void merge_groups_helper(unsigned long num_final_groups,
                      float blur_initial_sigma,
                      float blur_multiple, 
                      GROUP_T* segments,
                      MERGE_POLICY policy,
                      GaussianBlur<IMAGE_T>* gaussian);

    template <typename GROUP_T>
    void sequentialize_groups(GROUP_T* segments,
                              GROUP_T* group_map);

    template <typename GROUP_T>
    unsigned long sequentialize_groups_cpu(GROUP_T* segments,
                                           GROUP_T* group_map);

    template <typename GROUP_T>
    GROUP_T* allocate_array(unsigned long num_elements);

    template <typename GROUP_T>
    void free_array(GROUP_T*& arr);

};

#endif

---