Uniformization routine first stable build

Cpp/CMakeLists.txt

@@ -2,7 +2,7 @@
 # project specific logic here.
 #
 cmake_minimum_required (VERSION 3.20)
-
+# respect <PACKAGE>_ROOT variables in
 # Enable Hot Reload for MSVC compilers if supported.
 if (POLICY CMP0141)
   cmake_policy(SET CMP0141 NEW)
@@ -11,18 +11,31 @@ endif()
 
 set(CMAKE_CXX_STANDARD 20)
 
+# Non "standard" but common install prefixes
+list(APPEND CMAKE_SYSTEM_PREFIX_PATH
+#   /usr/X11R6
+#   /usr/pkg
+  /opt/
+)
+
+project ("main")
+
 # add libraries
 add_library(wfmLib "")
 add_library(AWGLib "")
+add_library(baslerLib "")
+add_library(imageProcLib "")
+add_library(uniformLib "")
 add_subdirectory(lib)
 include_directories(lib/driver_header)
 
 set(LIBS
     wfmLib
     AWGLib
+    baslerLib
+    imageProcLib
+    uniformLib
 )
 
-# main.exe
-project ("main")
+# executables
 add_subdirectory(src)
-target_link_libraries(main ${LIBS})

Cpp/lib/CMakeLists.txt

+add_subdirectory(devices)
 add_subdirectory(external)
 
 # find_package(Eigen3 3.4 REQUIRED)
@@ -13,16 +14,26 @@ target_sources(wfmLib
 # )
 # target_link_libraries(wfmLib helperLib)
 
-target_sources(AWGLib
+target_sources(imageProcLib
 	PRIVATE
-		AWG.cpp
+		image_process.cpp
 	PUBLIC
-		AWG.h
+		image_process.h
 )
-# add_library(driverLib STATIC IMPORTED)
-# set_target_properties(driverLib PROPERTIES
-#     IMPORTED_LOCATION "${CMAKE_CURRENT_LIST_DIR}/driver_header/spcm_win64_msvcpp.lib"
-#     INTERFACE_INCLUDE_DIRECTORIES "${CMAKE_CURRENT_LIST_DIR}/driver_header"
-# )
-
+find_package(OpenCV REQUIRED)
+include_directories(${OpenCV_INCLUDE_DIRS})
+target_link_libraries(imageProcLib ${OpenCV_LIBS})
 
+target_sources(uniformLib
+	PRIVATE
+		uniformization.cpp
+	PUBLIC
+		uniformization.h
+)
+set(UNILIBS
+	wfmLib
+	AWGLib
+	baslerLib
+	imageProcLib
+)
+target_link_libraries(uniformLib ${UNILIBS})

Cpp/lib/external/CMakeLists.txt

+# add_library(npy STATIC npy.hpp)
+# set_target_properties(npy PROPERTIES LINKER_LANGUAGE CXX)
\ No newline at end of file

Cpp/lib/image_process.cpp

+#include "image_process.h"
+#include <filesystem>
+#include <opencv4/opencv2/imgproc.hpp>
+#include <opencv4/opencv2/highgui.hpp>
+#include <opencv4/opencv2/imgcodecs.hpp>
+#include <iostream>
+#include "external/npy.hpp"
+
+// void TwzrImg::loadImage(
+//     std::vector<uint8_t> rawImg,
+//     uint width,
+//     uint height
+// ) {
+//     this->img = cv::Mat(height, width, CV_8U, rawImg.data());
+// }
+
+// void TwzrImg::loadImage(std::string filename) {
+//     this->img = cv::imread(filename, cv::IMREAD_GRAYSCALE);
+// }
+
+cv::Mat ImageProcessing::cvFromRaw(
+    uint8_t* buffer,
+    uint width,
+    uint height
+) {
+    return cv::Mat(height, width, CV_8U, buffer);
+}
+
+std::vector<std::vector<uint>> ImageProcessing::getPeaks(
+    const cv::Mat img,
+    double threshold,
+    std::string saveViewPath
+) {
+    double maxVal;
+    cv::minMaxIdx(img, NULL, &maxVal, NULL, NULL);
+    double threshVal = maxVal * threshold;
+
+    cv::Mat proc;
+    cv::GaussianBlur(img, proc, cv::Size(0,0), 3, 3);
+    if (threshold > 0) {
+        cv::threshold(proc, proc, threshVal, 255, cv::THRESH_BINARY);
+    } else {
+        cv::threshold(proc, proc, cv::THRESH_TRIANGLE, 255, cv::THRESH_BINARY);
+    }
+    cv::Mat labels, stats, centroids;
+    cv::connectedComponentsWithStats(
+        proc,
+        labels,
+        stats,
+        centroids,
+        4,
+        CV_32S,
+        cv::CCL_SAUF
+    );
+    std::vector<std::vector<uint>> peakLocs;
+    std::filesystem::path savePath(saveViewPath);
+    bool saveView = std::filesystem::exists(savePath) ? true : false;
+    cv::Mat imgCopy;
+    if (saveView) { img.copyTo(imgCopy); }
+    for (auto i = 1; i < centroids.rows; i++) {
+        cv::Point2d p(centroids.row(i));
+        std::vector<uint> row;
+        row.push_back(p.x);
+        row.push_back(p.y);
+        peakLocs.push_back(row);
+        if (saveView) { cv::circle(imgCopy, p, 3, 255); }
+    }
+    if (saveView) { 
+        // savePath = savePath / "figs";
+        // if (!std::filesystem::exists(savePath)) {
+        //     std::filesystem::create_directories(savePath);
+        // }
+        cv::imwrite(savePath / "peaks.png", imgCopy); 
+    } else { 
+        std::cout << savePath << " invalid path, "
+            << "write to file in ImageProcessing::getPeaks "
+            << "failed\n";
+    }
+    return peakLocs;
+}
+
+std::vector<cv::Mat> ImageProcessing::getSubImgs(
+    const cv::Mat img,
+    std::vector<std::vector<uint>> peaks,
+    std::string saveViewPath
+) {
+    std::filesystem::path savePath(saveViewPath);
+    bool saveView = std::filesystem::exists(savePath) ? true : false;
+    if (saveView) {
+        // savePath = savePath / "figs";
+        // if (!std::filesystem::exists(savePath)) {
+        //     std::filesystem::create_directories(savePath);
+        // }
+    } else {
+        std::cout << savePath << " invalid path, "
+            << "write to file in ImageProcessing::getSubImgs "
+            << "failed\n";
+    }
+    std::vector<cv::Mat> subImages;
+    auto w = 80;
+    auto h = 50;
+    for (int i = 0; auto peak : peaks) {
+        auto x = peak[0] - w / 2;
+        auto y = peak[1] - h / 2;
+        cv::Rect roi(x, y, w, h);
+        cv::Mat cropimg = img(roi);
+        subImages.push_back(cropimg);
+        if (saveView) {
+            std::string filename = "crop_" + std::to_string(i) + ".png";
+            cv::imwrite(savePath / filename, cropimg); 
+        }
+        ++i;
+    }
+    return subImages;
+}
+
+std::vector<double> ImageProcessing::getPowers(const std::vector<cv::Mat> imgs) {
+    std::vector<double> powers;
+    for (auto img : imgs) {
+        powers.push_back(cv::sum(img)[0]);
+    }
+    return powers;
+}

Cpp/lib/image_process.h

+#pragma once
+#include <opencv4/opencv2/imgproc.hpp>
+#include <opencv4/opencv2/core.hpp>
+#include <opencv4/opencv2/highgui.hpp>
+
+namespace ImageProcessing {
+
+    cv::Mat cvFromRaw(
+        uint8_t* buffer,
+        uint width,
+        uint height
+    );
+
+    std::vector<std::vector<uint>> getPeaks(
+        const cv::Mat img,
+        double threshold=-1,
+        std::string saveViewPath=""
+    );
+
+    std::vector<cv::Mat> getSubImgs(
+        const cv::Mat img,
+        std::vector<std::vector<uint>> peaks,
+        std::string saveViewPath=""
+    );
+    
+    std::vector<double> getPowers(const std::vector<cv::Mat> imgs);
+
+};
\ No newline at end of file

Cpp/lib/uniformization.cpp

+#include "uniformization.h"
+#include "external/npy.hpp"
+#include <filesystem>
+#include <ctime>
+
+Eigen::ArrayXd Uniformization::getWeightedMask(
+    Eigen::Ref<Eigen::ArrayXd> detProbe,
+    Eigen::Ref<Eigen::ArrayXd> initPower,
+    double polarizability,
+    double meanDepth
+) {
+    auto depths = meanDepth + (detProbe - detProbe.mean()) / polarizability;
+    Eigen::ArrayXd mask = depths / depths.mean() 
+        * initPower.mean() / initPower;
+    return mask;
+}
+
+void Uniformization::reloadAWG(AWG& awg, ArrayWaveform& wfm) {
+    auto wfmData = wfm.getStaticWaveform();
+    awg.cardStop();
+    awg.writeSeqModeSegment(
+        0,
+        wfmData.first,
+        wfmData.second
+    );
+    awg.cardRun();
+    awg.forceTrigger();
+}
+
+cv::Mat Uniformization::grabMulti(
+    BaslerCam& basler,
+    int numGrab
+) {
+    if (!basler.isGrabbing()) { basler.startGrabbing(); }
+    auto img = basler.getImage();
+    Eigen::Map<Eigen::Array<uint8_t, -1, 1>> imgEigen(img.data(), img.size());
+    auto width = basler.getWidth();
+    auto height = basler.getHeight();
+    for (int i = 0; i < numGrab - 1; i++) {
+        auto newImg = basler.getImage();
+        imgEigen += Eigen::Map<Eigen::Array<uint8_t, -1, 1>> (newImg.data(), newImg.size());
+    }
+    imgEigen /= numGrab;
+    cv::Mat imgCV = ImageProcessing::cvFromRaw(img.data(), width, height);
+    return imgCV;
+}
+
+void Uniformization::saveCSV(
+    std::string filename,
+    std::vector<double> data
+) {
+    std::ofstream savefile(filename);
+    if (savefile.is_open()) {
+        for (auto d : data) {
+            savefile << d << ",";
+        }
+        savefile.close();
+    } else {
+        std::cout << filename << " invalid, "
+            << "saving to file failed" << std::endl;
+    }
+}
+
+void Uniformization::run(
+    AWG& awg,
+    BaslerCam& basler,
+    ArrayWaveform& waveform,
+    const Uniformization::Params& config
+) {
+    std::filesystem::path probePath(config.probeScanPath);
+    if (!probePath.has_extension()) {
+        probePath /= "probe-det.npy";
+    }
+    if (!std::filesystem::exists(probePath)) {
+        std::cout << config.probeScanPath << " invalid, " 
+            << "returning" << std::endl;
+        return;
+    }
+    auto savePath = probePath.parent_path() / "uniformization";
+    if (!std::filesystem::exists(savePath)) {
+        std::filesystem::create_directory(savePath);
+    }
+
+    std::filesystem::path figSavePath = savePath / "figs";
+    if (!std::filesystem::exists(figSavePath)) {
+        std::filesystem::create_directory(figSavePath);
+    }
+    std::filesystem::path dataSavePath = savePath;
+
+    auto detProbe = npy::read_npy<double>(probePath).data;
+    std::reverse(detProbe.begin(), detProbe.end());
+
+    auto twzrImageInit = grabMulti(basler, config.numImgingAvg);
+    auto peaks = ImageProcessing::getPeaks(twzrImageInit, -1, figSavePath);
+    auto twzrCropsInit = ImageProcessing::getSubImgs(
+        twzrImageInit,
+        peaks,
+        figSavePath
+    );
+    if (twzrCropsInit.size() != config.numTweezer) {
+        std::cout
+            << "specified " << config.numTweezer << " tweezers, "
+            << twzrCropsInit.size() << " found, "
+            << "returning..." << std::endl;
+        return;
+    }
+    auto initPower = ImageProcessing::getPowers(twzrCropsInit);
+    auto weightedMask = getWeightedMask(
+        Eigen::Map<Eigen::ArrayXd> (detProbe.data(), detProbe.size()),
+        Eigen::Map<Eigen::ArrayXd> (initPower.data(), initPower.size()),
+        config.polarizability,
+        config.meanDepth
+    );
+
+    std::vector<double> powerErrTotal;
+    std::vector<double> finalPower = initPower;
+    for (int i = 0; i < config.maxLoop; i++) {
+        std::cout << "\r" << "running uniformization "
+            << "loop " << i+1 << "/" << config.maxLoop << std::flush;
+        auto twzrImage = grabMulti(basler, config.numImgingAvg);
+        auto twzrCrops = ImageProcessing::getSubImgs(
+            twzrImage,
+            peaks
+        );
+        auto powers = ImageProcessing::getPowers(twzrCrops);
+        Eigen::Map<Eigen::ArrayXd> powersEigen(powers.data(), powers.size());
+        auto powerErrs = (
+            weightedMask * powersEigen 
+            - (weightedMask * powersEigen).mean()
+        ) / (weightedMask * powersEigen).mean();
+        auto amps = waveform.getAmplitude();
+        Eigen::Map<Eigen::ArrayXd> ampsEigen(amps.data(), amps.size());
+        ampsEigen -= config.stepSize * powerErrs;
+        waveform.setAmplitude(amps);
+        reloadAWG(awg, waveform);
+
+        powerErrTotal.push_back(powerErrs.abs().sum());
+        if (i == config.maxLoop - 1) {
+            finalPower = powers;
+            std::cout << std::endl;
+        }
+    }
+    auto initPowerFile = dataSavePath / "initial_power.csv";
+    auto finalPowerFile = dataSavePath / "final_power.csv";
+    auto errorFile = dataSavePath / "errors.csv";
+    saveCSV(initPowerFile, initPower);
+    saveCSV(finalPowerFile, finalPower);
+    saveCSV(errorFile, powerErrTotal);
+}

Cpp/lib/uniformization.h

+#pragma once
+#include "waveform.h"
+#include "image_process.h"
+#include "devices/AWG.h"
+#include "devices/basler.h"
+
+namespace Uniformization {
+
+    struct Params {
+        std::string probeScanPath;
+        double polarizability;
+        double meanDepth;
+        double stepSize;
+        int maxLoop;
+        int numImgingAvg;
+        int numTweezer;
+    };
+
+    Eigen::ArrayXd getWeightedMask(
+        Eigen::Ref<Eigen::ArrayXd> detProbe,
+        Eigen::Ref<Eigen::ArrayXd> initPower,
+        double polarizability,
+        double meanDepth
+    );
+
+    void reloadAWG(
+        AWG& awg,
+        ArrayWaveform& wfm
+    );
+
+    cv::Mat grabMulti(
+        BaslerCam& basler,
+        int numGrab
+    );
+
+    void saveCSV(
+        std::string filename,
+        std::vector<double> data
+    );
+
+    void run(
+        AWG& awg,
+        BaslerCam& basler,
+        ArrayWaveform& waveform,
+        const Params& config
+    );
+
+};
\ No newline at end of file

Cpp/src/CMakeLists.txt

-add_executable(main example.cpp)
\ No newline at end of file
+add_executable(example example.cpp)
+target_link_libraries(example ${LIBS})
+
+add_executable(run run.cpp)
+target_link_libraries(run ${LIBS})
\ No newline at end of file