SceneSwitcher/plugins/video/preview-dialog.cpp

#include "preview-dialog.hpp"
#include "opencv-helpers.hpp"
#include "ui-helpers.hpp"

#include <QLayout>
#include <screenshot-helper.hpp>

namespace advss {

PreviewDialog::PreviewDialog(QWidget *parent)
	: QDialog(parent),
	  _scrollArea(new QScrollArea),
	  _imageLabel(new QLabel(this)),
	  _rubberBand(new QRubberBand(QRubberBand::Rectangle, this))
{
	setWindowTitle("Advanced Scene Switcher");
	setWindowFlags(windowFlags() | Qt::WindowMaximizeButtonHint |
		       Qt::WindowCloseButtonHint);

	_valueLabel = new QLabel("");
	_statusLabel = new QLabel(obs_module_text(
		"AdvSceneSwitcher.condition.video.showMatch.loading"));
	resize(parent->window()->size());

	// Center image
	auto wrapper = new QWidget();
	auto wrapperHLayout = new QHBoxLayout();
	wrapperHLayout->addStretch();
	wrapperHLayout->addWidget(_imageLabel);
	wrapperHLayout->addStretch();
	auto wrapperVLayout = new QVBoxLayout();
	wrapperVLayout->addStretch();
	wrapperVLayout->addLayout(wrapperHLayout);
	wrapperVLayout->addStretch();
	wrapper->setLayout(wrapperVLayout);

	_scrollArea->setBackgroundRole(QPalette::Dark);
	_scrollArea->setWidget(wrapper);
	_scrollArea->setWidgetResizable(true);
	QVBoxLayout *layout = new QVBoxLayout;
	layout->addWidget(_valueLabel);
	layout->addWidget(_statusLabel);
	layout->addWidget(_scrollArea);
	setLayout(layout);
}

void PreviewDialog::mousePressEvent(QMouseEvent *event)
{
	_selectingArea = true;
	if (_type == PreviewType::SHOW_MATCH) {
		return;
	}
	_origin = event->pos();
	_rubberBand->setGeometry(QRect(_origin, QSize()));
	_rubberBand->show();
}

void PreviewDialog::mouseMoveEvent(QMouseEvent *event)
{
	if (_type == PreviewType::SHOW_MATCH) {
		return;
	}
	_rubberBand->setGeometry(QRect(_origin, event->pos()).normalized());
}

void PreviewDialog::mouseReleaseEvent(QMouseEvent *)
{
	if (_type == PreviewType::SHOW_MATCH) {
		return;
	}
	auto selectionStart =
		_rubberBand->mapToGlobal(_rubberBand->rect().topLeft());
	QRect selectionArea(selectionStart, _rubberBand->size());

	auto imageStart =
		_imageLabel->mapToGlobal(_imageLabel->rect().topLeft());
	QRect imageArea(imageStart, _imageLabel->size());

	auto intersected = imageArea.intersected(selectionArea);
	QRect checksize(QPoint(intersected.topLeft() - imageStart),
			intersected.size());
	if (checksize.x() >= 0 && checksize.y() >= 0 && checksize.width() > 0 &&
	    checksize.height() > 0) {
		emit SelectionAreaChanged(checksize);
	}
	_selectingArea = false;
}

PreviewDialog::~PreviewDialog()
{
	Stop();
}

void PreviewDialog::ShowMatch()
{
	_type = PreviewType::SHOW_MATCH;
	_rubberBand->hide();
	_valueLabel->show();
	Start();
	_statusLabel->setText("");
}

void PreviewDialog::SelectArea()
{
	_selectingArea = false;
	_type = PreviewType::SELECT_AREA;
	_rubberBand->show();
	_valueLabel->hide();
	Start();
	DrawFrame();
	_statusLabel->setText(obs_module_text(
		"AdvSceneSwitcher.condition.video.selectArea.status"));
}

void PreviewDialog::Stop()
{
	_thread.quit();
	_thread.wait();
}

void PreviewDialog::closeEvent(QCloseEvent *)
{
	Stop();
}

void PreviewDialog::PatternMatchParametersChanged(
	const PatternMatchParameters &params)
{
	std::unique_lock<std::mutex> lock(_mtx);
	_patternMatchParams = params;
	_patternImageData = CreatePatternData(_patternMatchParams.image);
}

void PreviewDialog::ObjDetectParametersChanged(const ObjDetectParameters &params)
{
	std::unique_lock<std::mutex> lock(_mtx);
	_objDetectParams = std::make_shared<ObjDetectParameters>(params);
}

void PreviewDialog::OCRParametersChanged(const OCRParameters &params)
{
	std::unique_lock<std::mutex> lock(_mtx);
	_ocrParams = std::make_shared<OCRParameters>(params);
}

void PreviewDialog::VideoSelectionChanged(const VideoInput &video)
{
	std::unique_lock<std::mutex> lock(_mtx);
	_video = video;
}

void PreviewDialog::AreaParametersChanged(const AreaParameters &params)
{
	std::unique_lock<std::mutex> lock(_mtx);
	_areaParams = params;
}

void PreviewDialog::ConditionChanged(int cond)
{
	Stop();
	close();

	std::unique_lock<std::mutex> lock(_mtx);
	_condition = static_cast<VideoCondition>(cond);
}

void PreviewDialog::UpdateValue(double matchValue)
{
	std::string temp = obs_module_text(
		"AdvSceneSwitcher.condition.video.patternMatchValue");
	temp += "%.3f";
	_valueLabel->setText(QString::asprintf(temp.c_str(), matchValue));
}

void PreviewDialog::UpdateStatus(const QString &status)
{
	_statusLabel->setText(status);
}

void PreviewDialog::UpdateImage(const QPixmap &image)
{
	_imageLabel->setPixmap(image);
	_imageLabel->adjustSize();
	if (_type == PreviewType::SELECT_AREA && !_selectingArea) {
		DrawFrame();
	}
	emit NeedImage(_video, _type, _patternMatchParams, _patternImageData,
		       _objDetectParams, _ocrParams, _areaParams, _condition);
}

void PreviewDialog::Start()
{
	if (!_video.ValidSelection()) {
		DisplayMessage(obs_module_text(
			"AdvSceneSwitcher.condition.video.screenshotFail"));
		close();
		return;
	}

	if (_thread.isRunning()) {
		return;
	}

	auto worker = new PreviewImage(_mtx);
	worker->moveToThread(&_thread);
	connect(&_thread, &QThread::finished, worker, &QObject::deleteLater);
	connect(worker, &PreviewImage::ImageReady, this,
		&PreviewDialog::UpdateImage);
	connect(worker, &PreviewImage::StatusUpdate, this,
		&PreviewDialog::UpdateStatus);
	connect(worker, &PreviewImage::ValueUpdate, this,
		&PreviewDialog::UpdateValue);
	connect(this, &PreviewDialog::NeedImage, worker,
		&PreviewImage::CreateImage);
	_thread.start();

	emit NeedImage(_video, _type, _patternMatchParams, _patternImageData,
		       _objDetectParams, _ocrParams, _areaParams, _condition);
}

void PreviewDialog::DrawFrame()
{
	if (!_video.ValidSelection()) {
		return;
	}
	auto imageStart =
		_imageLabel->mapToGlobal(_imageLabel->rect().topLeft());
	auto windowStart = mapToGlobal(rect().topLeft());
	_rubberBand->resize(_areaParams.area.width, _areaParams.area.height);
	_rubberBand->move(
		_areaParams.area.x + (imageStart.x() - windowStart.x()),
		_areaParams.area.y + (imageStart.y() - windowStart.y()));
	_rubberBand->show();
}

static void markPatterns(cv::Mat &matchResult, QImage &image,
			 const cv::Mat &pattern)
{
	auto matchImg = QImageToMat(image);
	for (int row = 0; row < matchResult.rows; row++) {
		for (int col = 0; col < matchResult.cols; col++) {
			if (matchResult.at<float>(row, col) != 0.0) {
				rectangle(matchImg, {col, row},
					  cv::Point(col + pattern.cols,
						    row + pattern.rows),
					  cv::Scalar(255, 0, 0, 255), 2, 8, 0);
			}
		}
	}
}

static void markObjects(QImage &image, std::vector<cv::Rect> &objects)
{
	auto frame = QImageToMat(image);
	for (size_t i = 0; i < objects.size(); i++) {
		rectangle(frame, cv::Point(objects[i].x, objects[i].y),
			  cv::Point(objects[i].x + objects[i].width,
				    objects[i].y + objects[i].height),
			  cv::Scalar(255, 0, 0, 255), 2, 8, 0);
	}
}

PreviewImage::PreviewImage(std::mutex &mtx) : _mtx(mtx) {}

void PreviewImage::CreateImage(
	const VideoInput &video, PreviewType type,
	const PatternMatchParameters &patternMatchParams,
	const PatternImageData &patternImageData,
	std::shared_ptr<ObjDetectParameters> objDetectParams,
	std::shared_ptr<OCRParameters> ocrParams,
	const AreaParameters &areaParams, VideoCondition condition)
{
	auto source = obs_weak_source_get_source(video.GetVideo());
	QRect screenshotArea;
	if (areaParams.enable && type == PreviewType::SHOW_MATCH) {
		screenshotArea.setRect(areaParams.area.x, areaParams.area.y,
				       areaParams.area.width,
				       areaParams.area.height);
	}
	Screenshot screenshot(source, screenshotArea, true);
	obs_source_release(source);

	if (!video.ValidSelection() || !screenshot.IsDone()) {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.screenshotFail"));
		emit ImageReady(QPixmap());
		return;
	}

	if (screenshot.GetImage().width() == 0 ||
	    screenshot.GetImage().height() == 0) {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.screenshotEmpty"));
		emit ImageReady(QPixmap());
		return;
	}

	if (type == PreviewType::SHOW_MATCH) {
		std::unique_lock<std::mutex> lock(_mtx);
		// Will emit status label update
		MarkMatch(screenshot.GetImage(), patternMatchParams,
			  patternImageData, objDetectParams, ocrParams,
			  condition);
	} else {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.selectArea.status"));
	}
	emit ImageReady(QPixmap::fromImage(screenshot.GetImage()));
}

void PreviewImage::MarkMatch(
	QImage &screenshot, const PatternMatchParameters &patternMatchParams,
	const PatternImageData &patternImageData,
	std::shared_ptr<ObjDetectParameters> objDetectParams,
	std::shared_ptr<OCRParameters> ocrParams, VideoCondition condition)
{
	if (condition == VideoCondition::PATTERN) {
		MarkPatternMatch(screenshot, patternMatchParams,
				 patternImageData);
	} else if (condition == VideoCondition::OBJECT) {
		MarkObjectMatch(screenshot, objDetectParams);
	} else if (condition == VideoCondition::OCR) {
		MarkOCRMatch(screenshot, ocrParams);
	}
}

void PreviewImage::MarkPatternMatch(
	QImage &screenshot, const PatternMatchParameters &patternMatchParams,
	const PatternImageData &patternImageData)
{
	cv::Mat result;
	double matchValue = std::numeric_limits<double>::signaling_NaN();
	MatchPattern(screenshot, patternImageData, patternMatchParams.threshold,
		     result, &matchValue, patternMatchParams.useAlphaAsMask,
		     patternMatchParams.matchMode);
	emit ValueUpdate(matchValue);
	if (result.total() == 0 || countNonZero(result) == 0) {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.patternMatchFail"));
	} else if (result.cols == 1 && result.rows == 1) {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.patternMatchSuccessFullImage"));
	} else {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.patternMatchSuccess"));
		markPatterns(result, screenshot, patternImageData.rgbaPattern);
	}
}

void PreviewImage::MarkObjectMatch(
	QImage &screenshot,
	const std::shared_ptr<ObjDetectParameters> &objDetectParams)
{
	if (!objDetectParams) {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.objectMatchFail"));
		return;
	}
	auto model = objDetectParams->GetModel();
	if (!model) {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.objectMatchFail"));
		return;
	}
	auto objects = MatchObject(screenshot, *model,
				   objDetectParams->scaleFactor,
				   objDetectParams->minNeighbors,
				   objDetectParams->minSize.CV(),
				   objDetectParams->maxSize.CV());
	if (objects.empty()) {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.objectMatchFail"));
	} else {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.objectMatchSuccess"));
		markObjects(screenshot, objects);
	}
}

void PreviewImage::MarkOCRMatch(QImage &screenshot,
				const std::shared_ptr<OCRParameters> &ocrParams)
{
	if (!ocrParams) {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.ocrMatchFail"));
		return;
	}
	auto text = RunOCR(ocrParams->GetOCR(), screenshot, ocrParams->color,
			   ocrParams->colorThreshold);

	if (!text) {
		emit StatusUpdate(obs_module_text(
			"AdvSceneSwitcher.condition.video.ocrMatchFail"));
	}

	QString status(obs_module_text(
		"AdvSceneSwitcher.condition.video.ocrMatchSuccess"));
	emit StatusUpdate(status.arg(QString::fromStdString(*text)));
}

} // namespace advss