MultiThresholdPartitionStep.cpp

Go to the documentation of this file.

/*
 * MultiThresholdPartitionStep.cpp
 *
 *  Created on: Jan 17, 2017
 *      Author: mschefer
 */

#include <boost/random.hpp>

#include "SEImplementation/Measurement/MultithreadedMeasurement.h"
#include "SEImplementation/Partition/MultiThresholdPartitionStep.h"

#include "SEFramework/Image/VectorImage.h"
#include "SEFramework/Image/ProcessedImage.h"

#include "SEImplementation/Plugin/DetectionFramePixelValues/DetectionFramePixelValues.h"
#include "SEImplementation/Plugin/PixelBoundaries/PixelBoundaries.h"
#include "SEImplementation/Plugin/ShapeParameters/ShapeParameters.h"
#include "SEImplementation/Plugin/PixelCentroid/PixelCentroid.h"
#include "SEImplementation/Plugin/PeakValue/PeakValue.h"
#include "SEImplementation/Plugin/DetectionFrameImages/DetectionFrameImages.h"

#include "SEImplementation/Property/SourceId.h"

#include "SEImplementation/Segmentation/Lutz.h"

namespace SourceXtractor {

namespace {
  boost::random::mt19937 rng { ((unsigned int) time(NULL)) };
}

class MultiThresholdNode : public std::enable_shared_from_this<MultiThresholdNode> {
public:

  MultiThresholdNode(const std::vector<PixelCoordinate>& pixel_list, SeFloat threshold)
    : m_pixel_list(pixel_list), m_is_split(false), m_threshold(threshold) {
  }

  void addChild(std::shared_ptr<MultiThresholdNode> child) {
    m_children.push_back(child);
    child->m_parent = shared_from_this();
  }

  bool contains(const Lutz::PixelGroup& pixel_group) const {
    for (auto pixel : m_pixel_list) {
      if (pixel_group.pixel_list[0] == pixel) {
        return true;
      }
    }
    return false;
  }

  const std::vector<std::shared_ptr<MultiThresholdNode>>& getChildren() const {
    return m_children;
  }

  std::shared_ptr<MultiThresholdNode> getParent() const {
    return m_parent.lock();
  }

  double getTotalIntensity(VectorImage<DetectionImage::PixelType>& image, const PixelCoordinate& offset) const {
    DetectionImage::PixelType total_intensity = 0;
    for (const auto& pixel_coord : m_pixel_list) {
      total_intensity += (image.getValue(pixel_coord - offset) - m_threshold);
    }

    return total_intensity;
  }

  bool isSplit() const {
    return m_is_split;
  }

  void flagAsSplit() {
    m_is_split = true;
    auto parent = m_parent.lock();
    if (parent != nullptr) {
      parent->flagAsSplit();
    }
  }

  const std::vector<PixelCoordinate>& getPixels() const {
    return m_pixel_list;
  }

  void debugPrint() const {
    std::cout << "(" << m_pixel_list.size();

    for (auto& child : m_children) {
      std::cout << ", ";
      child->debugPrint();
    }

    std::cout << ")";
  }

  void addPixel(PixelCoordinate pixel) {
    m_pixel_list.emplace_back(pixel);
  }

  SeFloat getThreshold() const {
    return m_threshold;
  }

private:
  std::vector<PixelCoordinate> m_pixel_list;

  std::weak_ptr<MultiThresholdNode> m_parent;
  std::vector<std::shared_ptr<MultiThresholdNode>> m_children;

  bool m_is_split;

  SeFloat m_threshold;
};

std::vector<std::unique_ptr<SourceInterface>> MultiThresholdPartitionStep::partition(
    std::unique_ptr<SourceInterface> original_source) const {

  auto parent_source_id = original_source->getProperty<SourceId>().getSourceId();

  auto& detection_frame = original_source->getProperty<DetectionFrame>();

  auto& detection_frame_images = original_source->getProperty<DetectionFrameImages>();
  const auto labelling_image = detection_frame_images.getLockedImage(LayerFilteredImage);

  auto& pixel_boundaries = original_source->getProperty<PixelBoundaries>();

  auto& pixel_coords = original_source->getProperty<PixelCoordinateList>().getCoordinateList();

  auto offset = pixel_boundaries.getMin();
  auto thumbnail_image = VectorImage<DetectionImage::PixelType>::create(
      pixel_boundaries.getWidth(), pixel_boundaries.getHeight());
  thumbnail_image->fillValue(0);

  auto min_value = original_source->getProperty<PeakValue>().getMinValue() * .8;
  auto peak_value = original_source->getProperty<PeakValue>().getMaxValue();

  for (auto pixel_coord : pixel_coords) {
    auto value = labelling_image->getValue(pixel_coord);
    thumbnail_image->setValue(pixel_coord - offset, value);
  }

  auto root = std::make_shared<MultiThresholdNode>(pixel_coords, 0);

  std::list<std::shared_ptr<MultiThresholdNode>> active_nodes { root };
  std::list<std::shared_ptr<MultiThresholdNode>> junction_nodes;

  // Build the tree
  for (unsigned int i = 1; i < m_thresholds_nb; i++) {

    auto threshold = min_value * pow(peak_value / min_value, (double) i / m_thresholds_nb);
    auto subtracted_image = SubtractImage<DetectionImage::PixelType>::create(thumbnail_image, threshold);

    LutzList lutz;
    lutz.labelImage(*subtracted_image, offset);

    std::list<std::shared_ptr<MultiThresholdNode>> active_nodes_copy(active_nodes);
    for (auto& node : active_nodes_copy) {
      int nb_of_groups_inside = 0;
      for (auto& pixel_group : lutz.getGroups()) {
        if (pixel_group.pixel_list.size() >= m_min_deblend_area && node->contains(pixel_group)) {
          nb_of_groups_inside++;
        }
      }

      if (nb_of_groups_inside == 0) {
        active_nodes.remove(node);
      }

      if (nb_of_groups_inside > 1) {
        active_nodes.remove(node);
        junction_nodes.push_back(node);
        for (auto& pixel_group : lutz.getGroups()) {
          if (pixel_group.pixel_list.size() >= m_min_deblend_area && node->contains(pixel_group)) {
            auto new_node = std::make_shared<MultiThresholdNode>(pixel_group.pixel_list, threshold);
            node->addChild(new_node);
            active_nodes.push_back(new_node);
          }
        }
      }
    }
  }

  // Identify the sources
  double intensity_threshold = root->getTotalIntensity(*thumbnail_image, offset) * m_contrast;

  std::vector<std::shared_ptr<MultiThresholdNode>> source_nodes;
  while (!junction_nodes.empty()) {
    auto node = junction_nodes.back();
    junction_nodes.pop_back();

    int nb_of_children_above_threshold = 0;

    for (auto child : node->getChildren()) {
      if (child->getTotalIntensity(*thumbnail_image, offset) > intensity_threshold) {
        nb_of_children_above_threshold++;
      }
    }

    if (nb_of_children_above_threshold >= 2) {
      node->flagAsSplit();
      for (auto child : node->getChildren()) {
        if (child->getTotalIntensity(*thumbnail_image, offset) > intensity_threshold && !child->isSplit()) {
          source_nodes.push_back(child);
        }
      }
    }
  }

  std::vector<std::unique_ptr<SourceInterface>> sources;
  if (source_nodes.empty()) {
    sources.emplace_back(std::move(original_source)); // no split, just forward the source unchanged
    return sources;
  }

  for (auto source_node : source_nodes) {
    // remove pixels in the new sources from the image
    for (auto& pixel : source_node->getPixels()) {
      thumbnail_image->setValue(pixel - offset, 0);
    }

    auto new_source = m_source_factory->createSource();

    new_source->setProperty<PixelCoordinateList>(source_node->getPixels());
    new_source->setProperty<DetectionFrame>(detection_frame.getEncapsulatedFrame());

    sources.push_back(std::move(new_source));
  }

  auto new_sources = reassignPixels(sources, pixel_coords, thumbnail_image, source_nodes, offset);

  for (auto& new_source : new_sources) {
    new_source->setProperty<DetectionFrame>(detection_frame.getEncapsulatedFrame());
    new_source->setProperty<SourceId>(parent_source_id);
  }

  return new_sources;
}

std::vector<std::unique_ptr<SourceInterface>> MultiThresholdPartitionStep::reassignPixels(
    const std::vector<std::unique_ptr<SourceInterface>>& sources,
    const std::vector<PixelCoordinate>& pixel_coords,
    std::shared_ptr<VectorImage<DetectionImage::PixelType>> image,
    const std::vector<std::shared_ptr<MultiThresholdNode>>& source_nodes,
    const PixelCoordinate& offset
    ) const {

  std::vector<SeFloat> amplitudes;
  for (auto& source : sources) {
    auto& pixel_list = source->getProperty<PixelCoordinateList>().getCoordinateList();
    auto& shape_parameters = source->getProperty<ShapeParameters>();

    auto thresh = source->getProperty<PeakValue>().getMinValue();
    auto peak = source->getProperty<PeakValue>().getMaxValue();

    auto dist = pixel_list.size() / (2.0 * M_PI * shape_parameters.getAbcor() * shape_parameters.getEllipseA() * shape_parameters.getEllipseB());
    auto amp = dist < 70.0 ? thresh * expf(dist) : 4.0 * peak;

    // limit expansion ??
    if (amp > 4.0 * peak) {
      amp = 4.0 * peak;
    }

    amplitudes.push_back(amp);
  }

  for (auto pixel : pixel_coords) {
    if (image->getValue(pixel - offset) > 0) {
      SeFloat cumulated_probability = 0;
      std::vector<SeFloat> probabilities;

      SeFloat min_dist = std::numeric_limits<SeFloat>::max();
      std::shared_ptr<MultiThresholdNode> closest_source_node;

      int i = 0;
      for (auto& source : sources) {
        auto& shape_parameters = source->getProperty<ShapeParameters>();
        auto& pixel_centroid = source->getProperty<PixelCentroid>();

        auto dx = pixel.m_x - pixel_centroid.getCentroidX();
        auto dy = pixel.m_y - pixel_centroid.getCentroidY();

        auto dist = 0.5 * (shape_parameters.getEllipseCxx()*dx*dx +
            shape_parameters.getEllipseCyy()*dy*dy + shape_parameters.getEllipseCxy()*dx*dy) /
            shape_parameters.getAbcor();

        if (dist < min_dist) {
          min_dist = dist;
          closest_source_node = source_nodes[i];
        }

        cumulated_probability += dist < 70.0 ? amplitudes[i] * expf(-dist) : 0.0;

        probabilities.push_back(cumulated_probability);
        i++;
      }

      if (probabilities.back() > 1.0e-31) {
        auto drand = double(probabilities.back()) * boost::random::uniform_01<double>()(rng);

        unsigned int i=0;
        for (; i<probabilities.size() && drand >= probabilities[i]; i++);
        if (i < source_nodes.size()) {
          source_nodes[i]->addPixel(pixel);
        } else {
          std::cout << i << " oops " << drand << " " << probabilities.back() << std::endl;
        }

      } else {
        // select closest source
        closest_source_node->addPixel(pixel);
      }
    }
  }

  int total_pixels = 0;

  std::vector<std::unique_ptr<SourceInterface>> new_sources;
  for (auto source_node : source_nodes) {
    // remove pixels in the new sources from the image
    for (auto& pixel : source_node->getPixels()) {
      image->setValue(pixel - offset, 0);
    }

    auto new_source = m_source_factory->createSource();

    auto pixels = source_node->getPixels();
    total_pixels += pixels.size();

    new_source->setProperty<PixelCoordinateList>(pixels);

    new_sources.push_back(std::move(new_source));
  }

  return new_sources;
}

MultiThresholdPartitionStep::MultiThresholdPartitionStep(std::shared_ptr<SourceFactory> source_factory, SeFloat contrast,
    unsigned int thresholds_nb, unsigned int min_deblend_area) :
  m_source_factory(source_factory), m_contrast(contrast), m_thresholds_nb(thresholds_nb),
  m_min_deblend_area(min_deblend_area) {}

} // namespace