Jungfraujoch/image_analysis/MXAnalysisWithoutFPGA.cpp

// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "MXAnalysisWithoutFPGA.h"

#include "spot_finding/StrongPixelSet.h"
#include "../compression/JFJochDecompress.h"

#include "spot_finding/SpotUtils.h"
#include "bragg_prediction/BraggPredictionFactory.h"
#include "image_preprocessing/ImagePreprocessorCPU.h"

#include "azint/AzIntEngineCPU.h"
#include "roi/ROIIntegrationCPU.h"
#include "spot_finding/ImageSpotFinderCPU.h"
#include "spot_finding/AdaptiveSpotFinderCPU.h"
#include "bragg_integration/BraggIntegrationEngineCPU.h"
#ifdef JFJOCH_USE_CUDA
#include "azint/AzIntEngineGPU.h"
#include "roi/ROIIntegrationGPU.h"
#include "spot_finding/ImageSpotFinderGPU.h"
#include "spot_finding/AdaptiveSpotFinderGPU.h"
#include "image_preprocessing/ImagePreprocessorGPU.h"
#include "image_preprocessing/ImagePreprocessorBufferGPU.h"
#include "bragg_integration/BraggIntegrationEngineGPU.h"
#include "../common/CUDAWrapper.h"
#endif


MXAnalysisWithoutFPGA::MXAnalysisWithoutFPGA(const DiffractionExperiment &in_experiment,
                                             const AzimuthalIntegrationMapping &in_integration,
                                             const PixelMask &in_mask,
                                             IndexAndRefine &in_indexer,
                                             bool in_enable_fused_adaptive_gpu)
    : experiment(in_experiment),
      integration(in_integration),
      enable_fused_adaptive_gpu(in_enable_fused_adaptive_gpu),
      npixels(experiment.GetPixelsNum()),
      xpixels(experiment.GetXPixelsNum()),
      indexer(in_indexer),
      prediction(CreateBraggPrediction(experiment.IsRotationIndexing())),
      mask(in_mask),
      mask_resolution(experiment.GetPixelsNum(), false),
      mask_high_res(-1),
      mask_low_res(-1) {
#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() == 0) {
#endif
        preprocessor_buffer = std::make_unique<ImagePreprocessorBuffer>(experiment.GetPixelsNum());
        spotFinder = std::make_unique<ImageSpotFinderCPU>(experiment.GetXPixelsNum(), experiment.GetYPixelsNum());
        azint = std::make_unique<AzIntEngineCPU>(integration);
        preprocessor = std::make_unique<ImagePreprocessorCPU>(in_experiment, in_mask);
        bragg_engine = std::make_unique<BraggIntegrationEngineCPU>(in_experiment);
        if (experiment.ROI().size() >= 1)
            roi = std::make_unique<ROIIntegrationCPU>(experiment);
#ifdef JFJOCH_USE_CUDA
    } else {
        stream = std::make_shared<CudaStream>();
        preprocessor_buffer = std::make_unique<ImagePreprocessorBufferGPU>(experiment.GetPixelsNum());
        preprocessor = std::make_unique<ImagePreprocessorGPU>(in_experiment, in_mask, stream);
        spotFinder = std::make_unique<ImageSpotFinderGPU>(experiment.GetXPixelsNum(), experiment.GetYPixelsNum(), stream);
        azint = std::make_unique<AzIntEngineGPU>(integration, stream);
        bragg_engine = std::make_unique<BraggIntegrationEngineGPU>(in_experiment, stream);
        if (experiment.ROI().size() >= 1)
            roi = std::make_unique<ROIIntegrationGPU>(experiment, stream);
        if (enable_fused_adaptive_gpu) {
            // One GPU engine that computes the azimuthal profile and the adaptive spot mask in a single
            // image pass. fused_adaptive aliases it so Analyze() can lift the profile out of it.
            auto fused = std::make_unique<AdaptiveSpotFinderGPU>(integration, stream);
            fused_adaptive = fused.get();
            adaptiveSpotFinder = std::move(fused);
        }
    }
#endif
    if (!adaptiveSpotFinder)
        adaptiveSpotFinder = std::make_unique<AdaptiveSpotFinderCPU>(integration);
}

void MXAnalysisWithoutFPGA::Analyze(DataMessage &output,
                                    AzimuthalIntegrationProfile &profile,
                                    const SpotFindingSettings &spot_finding_settings) {
    if ((output.image.GetWidth() != xpixels)
        || (output.image.GetWidth() * output.image.GetHeight() != npixels))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch in pixel size");

    const auto compression_start_time = std::chrono::steady_clock::now();
    const uint8_t *image_ptr = output.image.GetUncompressedPtr(decompression_buffer);
    const auto compression_end_time = std::chrono::steady_clock::now();
    if (output.image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION)
        output.compression_time_s = std::chrono::duration<float>(compression_end_time - compression_start_time).count();

    const auto preprocessing_start_time = std::chrono::steady_clock::now();
    auto ret = preprocessor->Analyze(*preprocessor_buffer, image_ptr, output.image.GetMode());
    const auto preprocessing_end_time = std::chrono::steady_clock::now();
    output.preprocessing_time_s = std::chrono::duration<float>(preprocessing_end_time - preprocessing_start_time).count();

    // The fused GPU engine (rugnux offline, GPU, adaptive detection) produces the azimuthal profile as
    // a byproduct of spot finding, so the separate azint pass is skipped in that case and the profile is
    // lifted out of the finder below.
    const bool fused = enable_fused_adaptive_gpu && spot_finding_settings.enable
                    && spot_finding_settings.adaptive_threshold && fused_adaptive != nullptr;

    if (!fused) {
        const auto azint_start_time = std::chrono::steady_clock::now();
        azint->Run(*preprocessor_buffer, profile);
        const auto azint_end_time = std::chrono::steady_clock::now();
        output.azint_time_s = std::chrono::duration<float>(azint_end_time - azint_start_time).count();
    }

    if (roi)
        roi->Run(*preprocessor_buffer, output.roi);

    if (spot_finding_settings.enable) {
        // Update resolution mask
        if (mask_high_res != spot_finding_settings.high_resolution_limit
            || mask_low_res != spot_finding_settings.low_resolution_limit)
            UpdateMaskResolution(spot_finding_settings);

        const auto spot_finding_start_time = std::chrono::steady_clock::now();
        ImageSpotFinder &finder = spot_finding_settings.adaptive_threshold
                                      ? static_cast<ImageSpotFinder &>(*adaptiveSpotFinder)
                                      : *spotFinder;
        const std::vector<DiffractionSpot> spots = finder.Run(*preprocessor_buffer, spot_finding_settings, mask_resolution);
        SpotAnalyze(experiment, spot_finding_settings, spots, output);
        const auto spot_finding_end_time = std::chrono::steady_clock::now();
        output.spot_finding_time_s = std::chrono::duration<float>(spot_finding_end_time - spot_finding_start_time).count();

#ifdef JFJOCH_USE_CUDA
        if (fused) {
            // Lift the azimuthal profile the fused engine computed in the same pass; its azint cost is
            // folded into spot_finding_time_s above.
            profile.Clear(integration);
            profile += fused_adaptive->GetProfile();
            output.azint_time_s = 0.0f;
        }
#endif

        if (spot_finding_settings.indexing)
            indexer.ProcessImage(output, spot_finding_settings, *prediction,
                [this](const std::vector<Reflection> &predicted, size_t npredicted, int64_t image_number) {
                    return bragg_engine->Run(*preprocessor_buffer, predicted, npredicted, image_number);
                });
    }

    output.max_viable_pixel_value = ret.max_value;
    output.min_viable_pixel_value = ret.min_value;
    output.error_pixel_count = ret.error_pixel_count;
    output.saturated_pixel_count = ret.saturated_pixel_count;
    output.az_int_profile = profile.GetResult();
    output.az_int_profile_count = profile.GetPixelCount();
    output.az_int_profile_std = profile.GetStd();

    output.bkg_estimate = profile.GetBkgEstimate(integration.Settings());
    output.ice_ring_score = profile.GetIceRingScore(integration.Settings(),
                                                    spot_finding_settings.ice_ring_width_Q_recipA);
}

void MXAnalysisWithoutFPGA::RebuildROI() {
    if (experiment.ROI().empty()) {
        roi.reset();
        return;
    }
#ifdef JFJOCH_USE_CUDA
    if (stream) {
        roi = std::make_unique<ROIIntegrationGPU>(experiment, stream);
        return;
    }
#endif
    roi = std::make_unique<ROIIntegrationCPU>(experiment);
}

void MXAnalysisWithoutFPGA::AnalyzeROIOnly(DataMessage &output) {
    if ((output.image.GetWidth() != xpixels)
        || (output.image.GetWidth() * output.image.GetHeight() != npixels))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch in pixel size");

    const uint8_t *image_ptr = output.image.GetUncompressedPtr(decompression_buffer);
    preprocessor->Analyze(*preprocessor_buffer, image_ptr, output.image.GetMode());
    RunROIOnly(output);
}

void MXAnalysisWithoutFPGA::RunROIOnly(DataMessage &output) {
    output.roi.clear();
    if (roi)
        roi->Run(*preprocessor_buffer, output.roi);
}

void MXAnalysisWithoutFPGA::UpdateMaskResolution(const SpotFindingSettings &settings) {
    mask_low_res = settings.low_resolution_limit;
    mask_high_res = settings.high_resolution_limit;
    auto const &resolution_map = integration.Resolution();
    for (int i = 0; i < mask_resolution.size(); i++)
        mask_resolution[i] = (resolution_map[i] > mask_low_res) || (resolution_map[i] < mask_high_res);
}