Yolo_Label/mainwindow.cpp

476 lines
14 KiB
C++
Raw Normal View History

2018-10-29 07:26:06 +03:00
#include "mainwindow.h"
#include "ui_mainwindow.h"
#include <QFileDialog>
#include <QColorDialog>
#include <QKeyEvent>
#include <QDebug>
#include <QShortcut>
#include <QCollator>
2018-10-29 07:26:06 +03:00
#include <iomanip>
using std::cout;
using std::endl;
using std::ofstream;
using std::ifstream;
using std::string;
MainWindow::MainWindow(QWidget *parent) :
QMainWindow(parent),
ui(new Ui::MainWindow)
{
ui->setupUi(this);
2021-03-15 15:04:06 +03:00
connect(new QShortcut(QKeySequence(Qt::CTRL | Qt::Key_S), this), SIGNAL(activated()), this, SLOT(save_label_data()));
connect(new QShortcut(QKeySequence(Qt::CTRL | Qt::Key_C), this), SIGNAL(activated()), this, SLOT(clear_label_data()));
2018-10-29 07:26:06 +03:00
connect(new QShortcut(QKeySequence(Qt::Key_S), this), SIGNAL(activated()), this, SLOT(next_label()));
connect(new QShortcut(QKeySequence(Qt::Key_W), this), SIGNAL(activated()), this, SLOT(prev_label()));
connect(new QShortcut(QKeySequence(Qt::Key_A), this), SIGNAL(activated()), this, SLOT(prev_img()));
connect(new QShortcut(QKeySequence(Qt::Key_D), this), SIGNAL(activated()), this, SLOT(next_img()));
connect(new QShortcut(QKeySequence(Qt::Key_Space), this), SIGNAL(activated()), this, SLOT(next_img()));
2021-03-15 15:04:06 +03:00
connect(new QShortcut(QKeySequence(Qt::CTRL | Qt::Key_D), this), SIGNAL(activated()), this, SLOT(remove_img()));
connect(new QShortcut(QKeySequence(Qt::Key_Delete), this), SIGNAL(activated()), this, SLOT(remove_img()));
2018-10-29 07:26:06 +03:00
2018-12-14 07:44:18 +03:00
init_table_widget();
2018-10-29 07:26:06 +03:00
}
MainWindow::~MainWindow()
{
delete ui;
}
void MainWindow::on_pushButton_open_files_clicked()
{
bool bRetImgDir = false;
bool bRetObjFile = false;
2018-10-29 07:26:06 +03:00
2018-12-14 07:44:18 +03:00
open_img_dir(bRetImgDir);
2018-10-29 07:26:06 +03:00
if (!bRetImgDir) return ;
2018-10-29 07:26:06 +03:00
2018-12-14 07:44:18 +03:00
open_obj_file(bRetObjFile);
2018-10-29 07:26:06 +03:00
if (!bRetObjFile) return ;
2018-10-29 07:26:06 +03:00
init();
}
void MainWindow::on_pushButton_change_dir_clicked()
{
bool bRetImgDir = false;
open_img_dir(bRetImgDir);
if (!bRetImgDir) return ;
init();
}
2018-10-29 07:26:06 +03:00
void MainWindow::init()
{
ui->label_image->init();
2018-12-14 07:44:18 +03:00
init_button_event();
init_horizontal_slider();
2018-10-29 07:26:06 +03:00
set_label(0);
goto_img(0);
}
void MainWindow::set_label_progress(const int fileIndex)
{
QString strCurFileIndex = QString::number(fileIndex);
2018-12-14 05:01:40 +03:00
QString strEndFileIndex = QString::number(m_imgList.size() - 1);
2018-10-29 07:26:06 +03:00
ui->label_progress->setText(strCurFileIndex + " / " + strEndFileIndex);
}
void MainWindow::set_focused_file(const int fileIndex)
{
2018-12-14 05:01:40 +03:00
ui->label_file->setText("File: " + m_imgList.at(fileIndex));
2018-10-29 07:26:06 +03:00
}
2018-12-14 05:01:40 +03:00
void MainWindow::goto_img(const int fileIndex)
2018-10-29 07:26:06 +03:00
{
2018-12-14 07:44:18 +03:00
bool bIndexIsOutOfRange = (fileIndex < 0 || fileIndex > m_imgList.size() - 1);
if (bIndexIsOutOfRange) return;
2018-10-29 07:26:06 +03:00
2018-12-14 07:44:18 +03:00
m_imgIndex = fileIndex;
2018-10-29 07:26:06 +03:00
2018-12-14 07:44:18 +03:00
bool bImgOpened;
ui->label_image->openImage(m_imgList.at(m_imgIndex), bImgOpened);
2018-10-29 07:26:06 +03:00
2018-12-14 07:44:18 +03:00
ui->label_image->loadLabelData(get_labeling_data(m_imgList.at(m_imgIndex)));
ui->label_image->showImage();
2018-10-29 07:26:06 +03:00
2018-12-14 07:44:18 +03:00
set_label_progress(m_imgIndex);
set_focused_file(m_imgIndex);
2018-10-29 07:26:06 +03:00
2018-12-14 07:44:18 +03:00
//it blocks crash with slider change
2018-10-29 07:26:06 +03:00
ui->horizontalSlider_images->blockSignals(true);
2018-12-14 05:01:40 +03:00
ui->horizontalSlider_images->setValue(m_imgIndex);
2018-10-29 07:26:06 +03:00
ui->horizontalSlider_images->blockSignals(false);
}
2018-12-14 07:44:18 +03:00
void MainWindow::next_img(bool bSavePrev)
2018-10-29 07:26:06 +03:00
{
2018-12-14 07:44:18 +03:00
if(bSavePrev && ui->label_image->isOpened()) save_label_data();
goto_img(m_imgIndex + 1);
}
2018-10-29 07:26:06 +03:00
2018-12-14 07:44:18 +03:00
void MainWindow::prev_img(bool bSavePrev)
{
if(bSavePrev) save_label_data();
goto_img(m_imgIndex - 1);
2018-10-29 07:26:06 +03:00
}
void MainWindow::save_label_data()const
{
2018-12-14 05:01:40 +03:00
if(m_imgList.size() == 0) return;
2018-10-29 07:26:06 +03:00
2018-12-14 05:01:40 +03:00
QString qstrOutputLabelData = get_labeling_data(m_imgList.at(m_imgIndex));
2018-10-29 07:26:06 +03:00
ofstream fileOutputLabelData(qstrOutputLabelData.toStdString());
if(fileOutputLabelData.is_open())
{
for(int i = 0; i < ui->label_image->m_objBoundingBoxes.size(); i++)
{
ObjectLabelingBox objBox = ui->label_image->m_objBoundingBoxes[i];
2018-10-31 13:30:27 +03:00
if(ui->checkBox_cropping->isChecked())
{
QImage cropped = ui->label_image->crop(ui->label_image->cvtRelativeToAbsoluteRectInImage(objBox.box));
2018-10-31 13:30:27 +03:00
if(!cropped.isNull())
{
2018-12-14 05:01:40 +03:00
string strImgFile = m_imgList.at(m_imgIndex).toStdString();
strImgFile = strImgFile.substr( strImgFile.find_last_of('/') + 1,
strImgFile.find_last_of('.') - strImgFile.find_last_of('/') - 1);
2018-10-31 13:30:27 +03:00
cropped.save(QString().fromStdString(strImgFile) + "_cropped_" + QString::number(i) + ".png");
}
2018-10-31 13:30:27 +03:00
}
2018-11-01 11:37:46 +03:00
double midX = objBox.box.x() + objBox.box.width() / 2.;
double midY = objBox.box.y() + objBox.box.height() / 2.;
double width = objBox.box.width();
double height = objBox.box.height();
2018-10-29 07:26:06 +03:00
fileOutputLabelData << objBox.label;
fileOutputLabelData << " ";
fileOutputLabelData << std::fixed << std::setprecision(6) << midX;
fileOutputLabelData << " ";
fileOutputLabelData << std::fixed << std::setprecision(6) << midY;
fileOutputLabelData << " ";
fileOutputLabelData << std::fixed << std::setprecision(6) << width;
fileOutputLabelData << " ";
fileOutputLabelData << std::fixed << std::setprecision(6) << height << std::endl;
2018-10-29 07:26:06 +03:00
}
fileOutputLabelData.close();
ui->textEdit_log->setText(qstrOutputLabelData + " saved.");
}
}
void MainWindow::clear_label_data()
{
ui->label_image->m_objBoundingBoxes.clear();
ui->label_image->showImage();
}
2020-10-02 16:26:19 +03:00
void MainWindow::remove_img()
{
if(m_imgList.size() > 0) {
//remove a image
QFile::remove(m_imgList.at(m_imgIndex));
//remove a txt file
QString qstrOutputLabelData = get_labeling_data(m_imgList.at(m_imgIndex));
QFile::remove(qstrOutputLabelData);
m_imgList.removeAt(m_imgIndex);
if(m_imgList.size() == 0)
{
pjreddie_style_msgBox(QMessageBox::Information,"End", "In directory, there are not any image. program quit.");
QCoreApplication::quit();
}
else if( m_imgIndex == m_imgList.size())
{
m_imgIndex --;
}
goto_img(m_imgIndex);
}
}
2018-10-29 07:26:06 +03:00
void MainWindow::next_label()
{
2018-12-14 05:01:40 +03:00
set_label(m_objIndex + 1);
2018-10-29 07:26:06 +03:00
}
void MainWindow::prev_label()
{
2018-12-14 05:01:40 +03:00
set_label(m_objIndex - 1);
2018-10-29 07:26:06 +03:00
}
void MainWindow::load_label_list_data(QString qstrLabelListFile)
{
ifstream inputLabelListFile(qstrLabelListFile.toStdString());
if(inputLabelListFile.is_open())
{
2018-10-30 06:28:13 +03:00
for(int i = 0 ; i <= ui->tableWidget_label->rowCount(); i++)
ui->tableWidget_label->removeRow(ui->tableWidget_label->currentRow());
2018-12-14 05:01:40 +03:00
m_objList.clear();
ui->tableWidget_label->setRowCount(0);
2018-10-30 11:05:37 +03:00
ui->label_image->m_drawObjectBoxColor.clear();
2018-10-29 07:26:06 +03:00
string strLabel;
int fileIndex = 0;
while(getline(inputLabelListFile, strLabel))
2018-10-29 07:26:06 +03:00
{
int nRow = ui->tableWidget_label->rowCount();
2020-08-21 19:10:44 +03:00
2018-10-29 07:26:06 +03:00
QString qstrLabel = QString().fromStdString(strLabel);
QColor labelColor = label_img::BOX_COLORS[(fileIndex++)%10];
2018-12-14 05:01:40 +03:00
m_objList << qstrLabel;
2018-10-29 07:26:06 +03:00
ui->tableWidget_label->insertRow(nRow);
2018-10-29 11:37:18 +03:00
2018-10-29 07:26:06 +03:00
ui->tableWidget_label->setItem(nRow, 0, new QTableWidgetItem(qstrLabel));
2018-10-29 11:37:18 +03:00
ui->tableWidget_label->item(nRow, 0)->setFlags(ui->tableWidget_label->item(nRow, 0)->flags() ^ Qt::ItemIsEditable);
2018-10-29 07:26:06 +03:00
ui->tableWidget_label->setItem(nRow, 1, new QTableWidgetItem(QString().fromStdString("")));
2021-03-15 15:04:06 +03:00
ui->tableWidget_label->item(nRow, 1)->setBackground(labelColor);
2018-10-29 07:26:06 +03:00
ui->tableWidget_label->item(nRow, 1)->setFlags(ui->tableWidget_label->item(nRow, 1)->flags() ^ Qt::ItemIsEditable ^ Qt::ItemIsSelectable);
ui->label_image->m_drawObjectBoxColor.push_back(labelColor);
}
}
}
QString MainWindow::get_labeling_data(QString qstrImgFile)const
{
string strImgFile = qstrImgFile.toStdString();
string strLabelData = strImgFile.substr(0, strImgFile.find_last_of('.')) + ".txt";
return QString().fromStdString(strLabelData);
}
2018-12-14 05:01:40 +03:00
void MainWindow::set_label(const int labelIndex)
2018-10-29 07:26:06 +03:00
{
2018-12-14 05:01:40 +03:00
bool bIndexIsOutOfRange = (labelIndex < 0 || labelIndex > m_objList.size() - 1);
2018-10-29 07:26:06 +03:00
2018-12-14 05:01:40 +03:00
if(!bIndexIsOutOfRange)
2018-10-29 07:26:06 +03:00
{
2018-12-14 05:01:40 +03:00
m_objIndex = labelIndex;
ui->label_image->setFocusObjectLabel(m_objIndex);
ui->label_image->setFocusObjectName(m_objList.at(m_objIndex));
ui->tableWidget_label->setCurrentCell(m_objIndex, 0);
2018-10-29 07:26:06 +03:00
}
}
2018-12-14 07:44:18 +03:00
void MainWindow::set_label_color(const int index, const QColor color)
2018-10-29 07:26:06 +03:00
{
ui->label_image->m_drawObjectBoxColor.replace(index, color);
}
void MainWindow::pjreddie_style_msgBox(QMessageBox::Icon icon, QString title, QString content)
{
QMessageBox msgBox(icon, title, content, QMessageBox::Ok);
QFont font;
font.setBold(true);
msgBox.setFont(font);
msgBox.button(QMessageBox::Ok)->setFont(font);
msgBox.button(QMessageBox::Ok)->setStyleSheet("border-style: outset; border-width: 2px; border-color: rgb(0, 255, 0); color : rgb(0, 255, 0);");
msgBox.button(QMessageBox::Ok)->setFocusPolicy(Qt::ClickFocus);
2018-10-29 07:26:06 +03:00
msgBox.setStyleSheet("background-color : rgb(34, 0, 85); color : rgb(0, 255, 0);");
msgBox.exec();
}
2018-12-14 07:44:18 +03:00
void MainWindow::open_img_dir(bool& ret)
{
pjreddie_style_msgBox(QMessageBox::Information,"Help", "Step 1. Open Your Data Set Directory");
2021-03-15 15:04:06 +03:00
QString opened_dir;
if(m_imgDir.size() > 0) opened_dir = m_imgDir;
else opened_dir = QDir::currentPath();
QString imgDir = QFileDialog::getExistingDirectory(
nullptr,
tr("Open Dataset Directory"),
2021-03-15 15:04:06 +03:00
opened_dir,
QFileDialog::ShowDirsOnly);
QDir dir(imgDir);
QCollator collator;
collator.setNumericMode(true);
QStringList fileList = dir.entryList(
QStringList() << "*.jpg" << "*.JPG" << "*.png" << "*.bmp",
QDir::Files);
std::sort(fileList.begin(), fileList.end(), collator);
if(fileList.empty())
{
ret = false;
2018-12-14 07:44:18 +03:00
pjreddie_style_msgBox(QMessageBox::Critical,"Error", "This folder is empty");
}
else
{
2018-12-14 07:44:18 +03:00
ret = true;
m_imgDir = imgDir;
2018-12-14 05:01:40 +03:00
m_imgList = fileList;
2018-12-14 05:01:40 +03:00
for(QString& str: m_imgList)
str = m_imgDir + "/" + str;
}
}
2018-12-14 07:44:18 +03:00
void MainWindow::open_obj_file(bool& ret)
{
pjreddie_style_msgBox(QMessageBox::Information,"Help", "Step 2. Open Your Label List File(*.txt or *.names)");
2021-03-15 15:04:06 +03:00
QString opened_dir;
if(m_imgDir.size() > 0) opened_dir = m_imgDir;
else opened_dir = QDir::currentPath();
QString fileLabelList = QFileDialog::getOpenFileName(
nullptr,
tr("Open LabelList file"),
2021-03-15 15:04:06 +03:00
opened_dir,
tr("LabelList Files (*.txt *.names)"));
if(fileLabelList.size() == 0)
{
ret = false;
2018-12-14 07:44:18 +03:00
pjreddie_style_msgBox(QMessageBox::Critical,"Error", "LabelList file is not opened()");
}
else
{
ret = true;
2018-12-14 07:44:18 +03:00
load_label_list_data(fileLabelList);
}
}
2018-12-14 07:44:18 +03:00
void MainWindow::reupdate_img_list()
{
}
2018-10-29 07:26:06 +03:00
void MainWindow::wheelEvent(QWheelEvent *ev)
{
2021-03-15 15:04:06 +03:00
if(ev->angleDelta().y() > 0) // up Wheel
2018-10-29 07:26:06 +03:00
prev_img();
2021-03-15 15:04:06 +03:00
else if(ev->angleDelta().y() < 0) //down Wheel
2018-10-29 07:26:06 +03:00
next_img();
}
void MainWindow::on_pushButton_prev_clicked()
{
prev_img();
}
void MainWindow::on_pushButton_next_clicked()
{
next_img();
}
void MainWindow::keyPressEvent(QKeyEvent * event)
{
int nKey = event->key();
2018-10-29 07:26:06 +03:00
bool graveAccentKeyIsPressed = (nKey == Qt::Key_QuoteLeft);
bool numKeyIsPressed = (nKey >= Qt::Key_0 && nKey <= Qt::Key_9 );
2018-10-29 07:26:06 +03:00
if(graveAccentKeyIsPressed)
{
set_label(0);
2018-10-29 07:26:06 +03:00
}
else if(numKeyIsPressed)
{
int asciiToNumber = nKey - Qt::Key_0;
2018-12-14 05:01:40 +03:00
if(asciiToNumber < m_objList.size() )
2018-10-29 07:26:06 +03:00
{
set_label(asciiToNumber);
2018-10-29 07:26:06 +03:00
}
}
}
void MainWindow::on_pushButton_save_clicked()
{
save_label_data();
}
2018-12-14 07:44:18 +03:00
void MainWindow::on_pushButton_remove_clicked()
{
2020-10-02 16:26:19 +03:00
remove_img();
2018-12-14 07:44:18 +03:00
}
2018-10-29 07:26:06 +03:00
void MainWindow::on_tableWidget_label_cellDoubleClicked(int row, int column)
{
2018-12-14 05:01:40 +03:00
bool bColorAttributeClicked = (column == 1);
if(bColorAttributeClicked)
2018-10-29 07:26:06 +03:00
{
2018-12-18 08:58:30 +03:00
QColor color = QColorDialog::getColor(Qt::white,nullptr,"Set Label Color");
2018-10-29 07:26:06 +03:00
if(color.isValid())
{
set_label_color(row, color);
2021-03-15 15:04:06 +03:00
ui->tableWidget_label->item(row, 1)->setBackground(color);
2018-10-29 07:26:06 +03:00
}
set_label(row);
ui->label_image->showImage();
}
}
void MainWindow::on_tableWidget_label_cellClicked(int row, int column)
{
set_label(row);
}
void MainWindow::on_horizontalSlider_images_sliderMoved(int position)
{
goto_img(position);
}
2018-12-14 07:44:18 +03:00
void MainWindow::init_button_event()
{
ui->pushButton_change_dir->setEnabled(true);
2018-12-14 07:44:18 +03:00
ui->pushButton_next->setEnabled(true);
ui->pushButton_prev->setEnabled(true);
ui->pushButton_save->setEnabled(true);
ui->pushButton_remove->setEnabled(true);
}
void MainWindow::init_horizontal_slider()
2018-10-29 07:26:06 +03:00
{
2018-12-14 07:44:18 +03:00
ui->horizontalSlider_images->setEnabled(true);
ui->horizontalSlider_images->setRange(0, m_imgList.size() - 1);
ui->horizontalSlider_images->blockSignals(true);
ui->horizontalSlider_images->setValue(0);
ui->horizontalSlider_images->blockSignals(false);
2018-10-29 07:26:06 +03:00
}
2018-12-14 07:44:18 +03:00
void MainWindow::init_table_widget()
{
ui->tableWidget_label->horizontalHeader()->setVisible(true);
ui->tableWidget_label->horizontalHeader()->setStyleSheet("");
ui->tableWidget_label->horizontalHeader()->setStretchLastSection(true);
disconnect(ui->tableWidget_label->horizontalHeader(), SIGNAL(sectionPressed(int)),ui->tableWidget_label, SLOT(selectColumn(int)));
}