MainWindow.cpp

//
// Created by Naren Sadhwani on 03.04.24.
//

#include "MainWindow.h"
#include <QMenuBar>
#include <QToolBar>
#include <QSplitter>
#include <QGraphicsView>
#include <QLabel>
#include <QStatusBar>
#include <QFileDialog>
#include <QApplication>
#include <QMessageBox>
#include <QPixmap>
#include <QGraphicsPixmapItem>
#include <QTextStream>
#include <clocale>
#include <QTimer>

MainWindow::MainWindow(QWidget *parent) :QMainWindow(parent),
currentImage(nullptr), fileMenu(nullptr), tesseractAPI(nullptr){
    initUI();
}

void MainWindow::initUI() {
    this->resize(800, 600);
    // Setup Menubar
    fileMenu = menuBar()->addMenu("&File");
    fileMenuToolBar = addToolBar("File");

    //Main area
    auto *splitter = new QSplitter(Qt::Horizontal, this);
    imageScene = new QGraphicsScene(this);
    imageView =  new QGraphicsView(imageScene);

    splitter->addWidget(imageView);

    editor = new QTextEdit(this);
    splitter->addWidget(editor);

    QList<int> sizes = {400,400};
    splitter->setSizes(sizes);

    setCentralWidget(splitter);

    //Setup Status bar
    // setup status bar
    mainStatusBar = statusBar();
    mainStatusLabel = new QLabel(mainStatusBar);
    mainStatusBar->addPermanentWidget(mainStatusLabel);
    mainStatusLabel->setText("Application Information will be here!");

    // Detection Checkbox
    detectAreaCheckBox = new QCheckBox("Detect Text Areas?", this);

    createActions();
}

void MainWindow::createActions() {
    openAction = new QAction("&Open", this);
    fileMenu->addAction(openAction);
    saveImageAction = new QAction("Save &Image as", this);
    fileMenu->addAction(saveImageAction);
    saveTextAsAction = new QAction("Save &Text as", this);
    fileMenu->addAction(saveTextAsAction);
    exitAction = new QAction("&Quit", this);
    fileMenu->addAction(exitAction);
    ocrAction = new QAction("OCR", this);
    fileMenu->addAction(ocrAction);
    captureAction = new QAction("Capture Screen", this);


    fileMenuToolBar->addAction(openAction);
    fileMenuToolBar->addAction(saveImageAction);
    fileMenuToolBar->addAction(saveTextAsAction);
    fileMenuToolBar->addAction(exitAction);
    fileMenuToolBar->addAction(ocrAction);
    fileMenuToolBar->addWidget(detectAreaCheckBox);
    fileMenuToolBar->addAction(captureAction);

    // Connect slots to signals
    connect(exitAction, SIGNAL(triggered(bool)), QApplication::instance(), SLOT(quit()));
    connect(openAction, SIGNAL(triggered(bool)), this, SLOT(openImage()));
    connect(saveImageAction, SIGNAL(triggered(bool)), this, SLOT(saveImageAs()));
    connect(saveTextAsAction, SIGNAL(triggered(bool)), this, SLOT(saveTextAs()));
    connect(ocrAction, SIGNAL(triggered(bool)), this, SLOT(extractText()));
    connect(captureAction, SIGNAL(triggered(bool)), this, SLOT(captureScreen()));


    setupShortcuts();
}

void MainWindow::setupShortcuts() {
    //Open Image/Text
    QKeySequence shortcutOpen (Qt::Key_Super_L + Qt::Key_O);
    openAction->setShortcut(shortcutOpen);

    //Quit Application
    QKeySequence shortcutQuit (Qt::Key_Super_L +Qt::Key_Q);
    exitAction->setShortcut(shortcutQuit);
}


void MainWindow::openImage(){

    //qDebug() << "slot openImage is called";
    QFileDialog dialog(this);
    dialog.setWindowTitle("Open Image");
    dialog.setFileMode(QFileDialog::ExistingFile);
    dialog.setNameFilter(tr("Images (*.png *.bmp *.jpg)"));
    QStringList filePaths;
    if (dialog.exec()) {
        filePaths = dialog.selectedFiles();
        showImage(filePaths.at(0));
    }

}
void MainWindow::showImage(QPixmap image)
{
    imageScene->clear();
    imageView->resetTransform();
    currentImage = imageScene->addPixmap(image);
    imageScene->update();
    imageView->setSceneRect(image.rect());
}

void MainWindow::showImage(QString path)
{
    QPixmap image(path);
    showImage(image);
    currentImagePath = path;
    QString status = QString("%1, %2x%3, %4 Bytes").arg(path).arg(image.width())
            .arg(image.height()).arg(QFile(path).size());
    mainStatusLabel->setText(status);
}




void MainWindow::saveImageAs() {
    //Check if there is an image loaded, If not return
    if (currentImage == nullptr) {
        QMessageBox::information(this, "Error","No Image to save");
        return;
    }

    QFileDialog dialog(this);
    dialog.setWindowTitle("Save Image as: ");
    dialog.setFileMode(QFileDialog::AnyFile);
    dialog.setAcceptMode(QFileDialog::AcceptSave);
    dialog.setNameFilter(tr("*.jpeg *.jpg *.bmp *.svg *.png"));
    dialog.setDefaultSuffix("png");

    QStringList fileNames;

    if(dialog.exec()){
        fileNames = dialog.selectedFiles();
        if(QRegExp(".+\\.(png|bmp|jpg|jpeg|svg)").exactMatch(fileNames.at(0))) {
            currentImage->pixmap().save(fileNames.at(0));
        } else {
            QMessageBox::information(this, "Information", "Save error: bad format or filename.");
        }
    }

}

void MainWindow::saveTextAs() {
    QFileDialog dialog(this);
    dialog.setWindowTitle("Save Text as: ");
    dialog.setFileMode(QFileDialog::AnyFile);
    dialog.setAcceptMode(QFileDialog::AcceptSave);
    dialog.setNameFilter(tr("Text files (*.txt"));

    QStringList fileNames;
    if(dialog.exec()){
        fileNames = dialog.selectedFiles();
        if(QRegExp(".+\\.(txt)").exactMatch(fileNames.at(0))){
            QFile file(fileNames.at(0));
            if(!file.open(QIODevice::WriteOnly |QIODevice::Text)){
                QMessageBox::information(this, "Error", "Can't save text");
                return;
            }
            QTextStream out(&file);
            out << editor->toPlainText() <<"\n";
        } else {
            QMessageBox::information(this, "Error", "Save error. Bad format or filename.");
        }
    }


}

void MainWindow::extractText() {
    if(currentImage== nullptr){
        QMessageBox::information(this, "Error", "Please select an image.");
        return;
    }

    char* oldCtype = strdup(setlocale(LC_ALL, NULL));
    setlocale(LC_ALL, "C");
    if (tesseractAPI == nullptr) {
        tesseractAPI = new tesseract::TessBaseAPI();
        // Initialize tesseract-ocr with English, with specifying tessdata path
        if (tesseractAPI->Init(TESSDATA_PREFIX, "eng")) {
            QMessageBox::information(this, "Error", "Could not initialize tesseract.");
            return;
        }
    }
    QPixmap pixmap = currentImage->pixmap();
    QImage image = pixmap.toImage();
    image = image.convertToFormat(QImage::Format_RGB888);

    tesseractAPI->SetImage( image.bits(), image.width(), image.height(), 3, image.bytesPerLine());

    //Find text areas and output text on detected regions
    if (detectAreaCheckBox->checkState() == Qt::Checked){
        std::vector<cv::Rect> areas;
        cv::Mat newImage  = detectTextAreas(image, areas);
        showImage(newImage);
        for(cv::Rect &rect: areas){
            tesseractAPI->SetRectangle(rect.x, rect.y, rect.width, rect.height);
            char *outText = tesseractAPI->GetUTF8Text();
            editor->setPlainText(editor->toPlainText()+outText);
            delete [] outText;

        }
    } else {
        char *outText = tesseractAPI->GetUTF8Text();
        editor->setPlainText(outText);
        delete[] outText;
    }

    setlocale(LC_ALL, oldCtype);
    free(oldCtype);

}



void MainWindow::showImage(cv::Mat mat) {
    QImage image(mat.data, mat.cols, mat.rows, mat.step, QImage::Format_RGB888);
    QPixmap pixmap = QPixmap::fromImage(image);
    imageScene->clear();
    imageView->resetTransform();
    currentImage = imageScene->addPixmap(pixmap);
    imageScene->update();
    imageView->setSceneRect(pixmap.rect());
}

void MainWindow::decode(const cv::Mat& scores, const cv::Mat& geometry, float scoreThresh,
                        std::vector<cv::RotatedRect>& detections, std::vector<float>& confidences)
{
    CV_Assert(scores.dims == 4); CV_Assert(geometry.dims == 4);
    CV_Assert(scores.size[0] == 1); CV_Assert(scores.size[1] == 1);
    CV_Assert(geometry.size[0] == 1);  CV_Assert(geometry.size[1] == 5);
    CV_Assert(scores.size[2] == geometry.size[2]);
    CV_Assert(scores.size[3] == geometry.size[3]);

    detections.clear();
    const int height = scores.size[2];
    const int width = scores.size[3];
    for (int y = 0; y < height; ++y) {
        const auto* scoresData = scores.ptr<float>(0, 0, y);
        const auto* x0_data = geometry.ptr<float>(0, 0, y);
        const auto* x1_data = geometry.ptr<float>(0, 1, y);
        const auto* x2_data = geometry.ptr<float>(0, 2, y);
        const auto* x3_data = geometry.ptr<float>(0, 3, y);
        const auto* anglesData = geometry.ptr<float>(0, 4, y);
        for (int x = 0; x < width; ++x) {
            float score = scoresData[x];
            if (score < scoreThresh)
                continue;

            // Decode a prediction.
            // Multiple by 4 because feature maps are 4 time less than input image.
            float offsetX = x * 4.0f, offsetY = y * 4.0f;
            float angle = anglesData[x];
            float cosA = std::cos(angle);
            float sinA = std::sin(angle);
            float h = x0_data[x] + x2_data[x];
            float w = x1_data[x] + x3_data[x];

            cv::Point2f offset(offsetX + cosA * x1_data[x] + sinA * x2_data[x],
                               offsetY - sinA * x1_data[x] + cosA * x2_data[x]);
            cv::Point2f p1 = cv::Point2f(-sinA * h, -cosA * h) + offset;
            cv::Point2f p3 = cv::Point2f(-cosA * w, sinA * w) + offset;
            cv::RotatedRect r(0.5f * (p1 + p3), cv::Size2f(w, h), -angle * 180.0f / (float)CV_PI);
            detections.push_back(r);
            confidences.push_back(score);
        }
    }
}

cv::Mat MainWindow::detectTextAreas(QImage &image, std::vector<cv::Rect> &areas) {
    float confidenceThreshold = 0.5;
    float nmsThreshold = 0.4;
    int inputWidth = 320;
    int inputHeight = 320;
    std::string model = "/Users/mcking/CLionProjects/OCR/EAST_Model/frozen_east_text_detection.pb";
    //Load NN
    if(net.empty()){
        net = cv::dnn::readNet(model);
    }

    std::vector<cv::Mat> outs;
    std::vector<std::string> layerNames(2);
    layerNames[0] = "feature_fusion/Conv_7/Sigmoid";
    layerNames[1] = "feature_fusion/concat_3";

    cv::Mat frame = cv::Mat(image.height(), image.width(), CV_8UC3, image.bits(), image.bytesPerLine());

    cv::Mat blob;

    cv::dnn::blobFromImage(frame, blob, 1.0, cv::Size(inputWidth,inputHeight), cv::Scalar(123.68, 116.78,103.94), true, false);

    net.setInput(blob);
    net.forward(outs, layerNames);

    cv::Mat scores = outs[0];
    cv::Mat geometry = outs[1];

    std::vector<cv::RotatedRect> boxes;
    std::vector<float> confidences;
    decode(scores, geometry, confidenceThreshold, boxes, confidences);

    std::vector<int> indices;
    cv::dnn::NMSBoxes(boxes, confidences, confidenceThreshold, nmsThreshold, indices);

    //Render detections
    cv::Point2f ratio((float)frame.cols / inputWidth, (float)frame.rows / inputHeight);
    cv::Scalar green = cv::Scalar(0, 255, 0);

    for (size_t i = 0; i < indices.size(); ++i) {
        cv::RotatedRect& box = boxes[indices[i]];
        cv::Rect area = box.boundingRect();
        area.x *= ratio.x;
        area.width *= ratio.x;
        area.y *= ratio.y;
        area.height *= ratio.y;
        areas.push_back(area);
        cv::rectangle(frame, area, green, 1);
        QString index = QString("%1").arg(i);
        cv::putText(
                frame, index.toStdString(), cv::Point2f(area.x, area.y - 2),
                cv::FONT_HERSHEY_SIMPLEX, 0.5, green, 1
        );
    }
    return frame;
}

MainWindow::~MainWindow() {
    if(tesseractAPI != nullptr) {
        tesseractAPI->End();
        delete tesseractAPI;        }
}

void MainWindow::captureScreen() {
    this->setWindowState(this->windowState() | Qt::WindowMinimized);
    QTimer::singleShot(500,this, SLOT(startCapture()));
}

void MainWindow::startCapture() {
    auto *cap = new ScreenCapture(this);
    cap->show();
    cap->activateWindow();
};