Skip to content

Commit

Permalink
SE-2767 fix saving
Browse files Browse the repository at this point in the history
  • Loading branch information
guy-martin committed May 23, 2024
1 parent 4e84315 commit 434cb96
Show file tree
Hide file tree
Showing 5 changed files with 129 additions and 87 deletions.
18 changes: 16 additions & 2 deletions stereo_viewer/inc/io/data_thread.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include <QWaitCondition>
#include <cstdint>
#include <QVector>
#include <QPair>
#include <opencv2/opencv.hpp>
#include <opencv2/core.hpp>
#include "inc/bottlenose_chunk_parser.hpp"
Expand All @@ -33,6 +34,15 @@
#define __IO_DATA_THREAD_HPP__

namespace labforge::io {
enum ImageDataType {
IMTYPE_IO, ///< single image
IMTYPE_DO, ///< disparity alone
IMTYPE_LR, ///< left+right images
IMTYPE_LD, ///< left+disparity images
IMTYPE_DR, ///< disparity+right images
IMTYPE_DC, ///< disparity+confidence images
};

struct ImageData
{
uint64_t timestamp;
Expand All @@ -42,8 +52,10 @@ struct ImageData
cv::Mat disparity;
int32_t min_disparity;
pointcloud_t pc;
ImageDataType imtype;
};


class DataThread : public QThread
{
Q_OBJECT
Expand All @@ -56,6 +68,7 @@ class DataThread : public QThread
const QImage &right, QString format,
const uint16_t *raw, int32_t,
const pointcloud_t &pc);
void setImageDataType(ImageDataType imtype);
bool setFolder(QString new_folder);
void setStereoDisparity(bool is_stereo, bool is_disparity);
void stop();
Expand All @@ -82,10 +95,11 @@ class DataThread : public QThread
QString m_left_fname;
QString m_right_fname;
QString m_disparity_fname;
QString m_conf_fname;
QString m_pc_fname;
bool m_stereo;

bool m_abort;
bool m_disparity;
ImageDataType m_imtype;
cv::Mat m_matQ;
};
}
Expand Down
3 changes: 2 additions & 1 deletion stereo_viewer/inc/ui/MainWindow.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
#include <PvStreamGEV.h>
#include <chrono>
#include <QProgressDialog>
#include <QPair>
#include <QtWidgets>

#include "ui_stereo_viewer.h"
Expand Down Expand Up @@ -57,7 +58,7 @@ public Q_SLOTS:
void handleStereoData(bool is_disparity);
void handleMonoData(bool is_disparity);
void handleError(QString msg);
void newData(uint64_t timestamp, QImage &left, QImage &right, bool stereo,
void newData(uint64_t timestamp, QImage &left, QImage &right, QPair<QString, QString> &label,
bool disparity, uint16_t *raw_disparity, int32_t min_disparity, const pointcloud_t &pc);
void onFolderSelect();
void handleSave();
Expand Down
7 changes: 4 additions & 3 deletions stereo_viewer/src/gev/pipeline.cc
Original file line number Diff line number Diff line change
Expand Up @@ -250,13 +250,14 @@ void Pipeline::run() {
QMutexLocker l(&m_image_lock);
img0 = lBuffer->GetImage();
m_pixformat->GetValue( pixformat );
int cv_pixformat = CV_16UC1;
//int cv_pixformat = CV_16UC1;
is_disparity = true;
int cv_pixformat = (img0->GetPixelType() == PvPixelYUV422_8)? CV_8UC2: CV_16UC1;

if(strcmp(pixformat.GetAscii(), "YUV422_8") == 0){
/*if(strcmp(pixformat.GetAscii(), "YUV422_8") == 0){
cv_pixformat = CV_8UC2;
is_disparity = false;
}
}*/

m_images.enqueue({new Mat(img0->GetHeight(), img0->GetWidth(), cv_pixformat, img0->GetDataPointer()),
new Mat(), timestamp, static_cast<int32_t>(minDisparity), pointcloud});
Expand Down
129 changes: 69 additions & 60 deletions stereo_viewer/src/io/data_thread.cc
Original file line number Diff line number Diff line change
Expand Up @@ -31,16 +31,14 @@ using namespace labforge::io;
using namespace std;

DataThread::DataThread(QObject *parent)
: QThread(parent), m_abort(false)
: QThread(parent), m_abort(false), m_imtype(IMTYPE_LR)
{
m_folder = "";
m_left_subfolder = "cam0";
m_right_subfolder = "cam1";
m_disparity_subfolder = "disparity";
m_pc_subfolder = "pc";
m_frame_counter = 0;
m_stereo = true;
m_disparity = false;
}

DataThread::~DataThread()
Expand All @@ -63,7 +61,8 @@ void DataThread::process(uint64_t timestamp, const QImage &left_image,
if(raw != nullptr){
dmat = cv::Mat(left_image.height(),left_image.width(), CV_16UC1, (uint16_t *)raw);
}
m_queue.enqueue({timestamp, left_image, right_image, format, dmat, min_disparity, pc});

m_queue.enqueue({timestamp, left_image, right_image, format, dmat, min_disparity, pc, m_imtype});

if (!isRunning()) {
start(HighPriority);
Expand All @@ -72,6 +71,10 @@ void DataThread::process(uint64_t timestamp, const QImage &left_image,
}
}

void DataThread::setImageDataType(ImageDataType imtype){
m_imtype = imtype;
}

bool getFilename(QString &fname, const QString &new_folder, const QString &subfolder, QString file_prefix){
QDir qdir(new_folder);
QString subdir_path = qdir.filePath(subfolder);
Expand All @@ -93,31 +96,27 @@ bool DataThread::setFolder(QString new_folder){
m_frame_counter = 0;
}

if(m_stereo){
if(m_disparity){
status = getFilename(m_left_fname, m_folder, m_left_subfolder, "left_");
status = status && getFilename(m_disparity_fname, m_folder, m_disparity_subfolder, "disparity_");
}else{
status = getFilename(m_left_fname, m_folder, m_left_subfolder, "left_");
status = status && getFilename(m_right_fname, m_folder, m_right_subfolder, "right_");
}
status = status && getFilename(m_pc_fname, m_folder, m_pc_subfolder, "spc_");
}else{
if(m_disparity){
status = getFilename(m_disparity_fname, m_folder, m_disparity_subfolder, "disparity_");
} else {
status = getFilename(m_left_fname, m_folder, m_left_subfolder, "mono_");
}
if(m_imtype == IMTYPE_IO){
status = getFilename(m_left_fname, m_folder, m_left_subfolder, "mono_");
} else if (m_imtype == IMTYPE_DO){
status = status && getFilename(m_disparity_fname, m_folder, m_disparity_subfolder, "disparity_");
} else if (m_imtype == IMTYPE_LR){
status = status && getFilename(m_left_fname, m_folder, m_left_subfolder, "left_");
status = status && getFilename(m_right_fname, m_folder, m_right_subfolder, "right_");
} else if (m_imtype == IMTYPE_LD){
status = status && getFilename(m_left_fname, m_folder, m_left_subfolder, "left_");
status = status && getFilename(m_disparity_fname, m_folder, m_disparity_subfolder, "disparity_");
} else if (m_imtype == IMTYPE_DR){
status = status && getFilename(m_right_fname, m_folder, m_right_subfolder, "right_");
status = status && getFilename(m_disparity_fname, m_folder, m_disparity_subfolder, "disparity_");
} else if (m_imtype == IMTYPE_DC){
status = status && getFilename(m_disparity_fname, m_folder, m_disparity_subfolder, "disparity_");
status = status && getFilename(m_conf_fname, m_folder, m_disparity_subfolder, "conf_");
}

return status;
}

void DataThread::setStereoDisparity(bool is_stereo, bool is_disparity){
m_stereo = is_stereo;
m_disparity = is_disparity;
}

void DataThread::stop(){
QMutexLocker locker(&m_mutex);
m_queue.clear();
Expand Down Expand Up @@ -149,7 +148,7 @@ static uint32_t countNaN(const pointcloud_t& pointCloud){
return count;
}

static void saveColoredPLYFile(const cv::Mat& pointCloud, QImage &image, const QString& filename) {
static void saveColoredPLYFile(const cv::Mat& pointCloud, const QImage &image, const QString& filename) {
QFile file(filename);

if (!file.open(QIODevice::WriteOnly | QIODevice::Text)) {
Expand Down Expand Up @@ -247,6 +246,24 @@ static void filterPointCloud(cv::Mat &pc, const cv::Mat &disp){
}
}

static void saveProjected3D(const cv::Mat &disparity, QImage &qimage,
int32_t min_disparity, cv::Mat &matQ, QString &dfname){
if(disparity.empty()) return;
if(dfname.isEmpty()) return;

cv::Mat pc(qimage.height(), qimage.width(), CV_32FC3);
cv::Mat dispf32;

disparity.convertTo(dispf32, CV_32FC1, (1./255.0), 0);
dispf32 += min_disparity;

if(!matQ.empty()){
cv::reprojectImageTo3D(dispf32, pc, matQ, false, CV_32F);
filterPointCloud(pc, disparity);
saveColoredPLYFile(pc, qimage, dfname);
}
}

void DataThread::run() {
while(!m_abort) {
m_mutex.lock();
Expand All @@ -265,41 +282,33 @@ void DataThread::run() {
QString suffix = padded_cntr + "_" + QString::number(imdata.timestamp) + "." + ext.toLower();
int32_t quality = (ext == "JPG") ? 90 : -1;

if (m_stereo){
if(m_disparity){
imdata.left.save(m_left_fname + suffix, ext.toStdString().c_str(), quality);
imdata.right.save(m_disparity_fname + suffix, ext.toStdString().c_str(), quality);

if(!imdata.disparity.empty()){
cv::Mat pc(imdata.left.height(), imdata.left.width(), CV_32FC3);
cv::Mat dispf32;

imdata.disparity.convertTo(dispf32, CV_32FC1, (1./255.0), 0);
//dispf32.setTo(-imdata.min_disparity, dispf32>255); ///< unmatched points to inf.
//dispf32.setTo(-imdata.min_disparity, dispf32==0);
dispf32 += imdata.min_disparity;

QString fname = m_disparity_fname + suffix.replace(ext.toLower(), "ply");
if(!m_matQ.empty()){
cv::reprojectImageTo3D(dispf32, pc, m_matQ, false, CV_32F);
filterPointCloud(pc, imdata.disparity);
saveColoredPLYFile(pc, imdata.left, fname);
}
}
} else {
imdata.left.save(m_left_fname + suffix, ext.toStdString().c_str(), quality);
imdata.right.save(m_right_fname + suffix, ext.toStdString().c_str(), quality);
}
if(imdata.pc.size() > 0){
QString fname = m_pc_fname + suffix.replace(ext.toLower(), "ply");
saveColoredPLYFile(imdata.pc, imdata.left, fname);
}
} else {
if(m_disparity){
imdata.left.save(m_disparity_fname + suffix, ext.toStdString().c_str(), quality);
} else {
imdata.left.save(m_left_fname + suffix, ext.toStdString().c_str(), quality);
}
if(imdata.imtype == IMTYPE_IO){
imdata.left.save(m_left_fname + suffix, ext.toStdString().c_str(), quality);
} else if (imdata.imtype == IMTYPE_DO){
imdata.left.save(m_disparity_fname + suffix, ext.toStdString().c_str(), quality);
} else if (imdata.imtype == IMTYPE_LR){
imdata.left.save(m_left_fname + suffix, ext.toStdString().c_str(), quality);
imdata.right.save(m_right_fname + suffix, ext.toStdString().c_str(), quality);
} else if (imdata.imtype == IMTYPE_LD){
imdata.left.save(m_left_fname + suffix, ext.toStdString().c_str(), quality);
imdata.right.save(m_disparity_fname + suffix, ext.toStdString().c_str(), quality);
QString fname = m_disparity_fname + suffix.replace(ext.toLower(), "ply");
saveProjected3D(imdata.disparity, imdata.left, imdata.min_disparity, m_matQ, fname);
} else if (imdata.imtype == IMTYPE_DR){
imdata.left.save(m_disparity_fname + suffix, ext.toStdString().c_str(), quality);
imdata.right.save(m_right_fname + suffix, ext.toStdString().c_str(), quality);
//QString fname = m_disparity_fname + suffix.replace(ext.toLower(), "ply");
//saveProjected3D(imdata.disparity, imdata.rig, imdata.min_disparity, m_matQ, fname);
} else if (imdata.imtype == IMTYPE_DC){
imdata.left.save(m_disparity_fname + suffix, ext.toStdString().c_str(), quality);
imdata.right.save(m_conf_fname + suffix, ext.toStdString().c_str(), quality);
//QString fname = m_disparity_fname + suffix.replace(ext.toLower(), "ply");
//saveProjected3D(imdata, m_matQ, fname);
}
if((imdata.pc.size() > 0) && (imdata.imtype == IMTYPE_LR)){
getFilename(m_pc_fname, m_folder, m_pc_subfolder, "spc_");
QString fname = m_pc_fname + suffix.replace(ext.toLower(), "ply");
saveColoredPLYFile(imdata.pc, imdata.left, fname);
}

m_frame_counter += 1;
Expand Down
Loading

0 comments on commit 434cb96

Please sign in to comment.