114 lines
3.3 KiB
Python
114 lines
3.3 KiB
Python
try:
|
|
from PyQt5.QtGui import *
|
|
from PyQt5.QtCore import *
|
|
from PyQt5.QtWidgets import *
|
|
except ImportError:
|
|
from PyQt4.QtGui import *
|
|
from PyQt4.QtCore import *
|
|
|
|
from libs.utils import newIcon
|
|
|
|
import time
|
|
import datetime
|
|
import json
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
BB = QDialogButtonBox
|
|
|
|
class DataPartitionDialog(QDialog):
|
|
def __init__(self, parent=None):
|
|
super().__init__()
|
|
self.parnet = parent
|
|
self.title = 'DATA PARTITION'
|
|
|
|
self.train_ratio = 70
|
|
self.val_ratio = 15
|
|
self.test_ratio = 15
|
|
|
|
self.initUI()
|
|
|
|
def initUI(self):
|
|
self.setWindowTitle(self.title)
|
|
self.setWindowModality(Qt.ApplicationModal)
|
|
|
|
self.flag_accept = True
|
|
|
|
if self.parnet.lang == 'ch':
|
|
msg = "导出JSON前请保存所有图像的标注且关闭EXCEL!"
|
|
else:
|
|
msg = "Please save all the annotations and close the EXCEL before exporting JSON!"
|
|
|
|
info_msg = QLabel(msg, self)
|
|
info_msg.setWordWrap(True)
|
|
info_msg.setStyleSheet("color: red")
|
|
info_msg.setFont(QFont('Arial', 12))
|
|
|
|
train_lbl = QLabel('Train split: ', self)
|
|
train_lbl.setFont(QFont('Arial', 15))
|
|
val_lbl = QLabel('Valid split: ', self)
|
|
val_lbl.setFont(QFont('Arial', 15))
|
|
test_lbl = QLabel('Test split: ', self)
|
|
test_lbl.setFont(QFont('Arial', 15))
|
|
|
|
self.train_input = QLineEdit(self)
|
|
self.train_input.setFont(QFont('Arial', 15))
|
|
self.val_input = QLineEdit(self)
|
|
self.val_input.setFont(QFont('Arial', 15))
|
|
self.test_input = QLineEdit(self)
|
|
self.test_input.setFont(QFont('Arial', 15))
|
|
|
|
self.train_input.setText(str(self.train_ratio))
|
|
self.val_input.setText(str(self.val_ratio))
|
|
self.test_input.setText(str(self.test_ratio))
|
|
|
|
validator = QIntValidator(0, 100)
|
|
self.train_input.setValidator(validator)
|
|
self.val_input.setValidator(validator)
|
|
self.test_input.setValidator(validator)
|
|
|
|
gridlayout = QGridLayout()
|
|
gridlayout.addWidget(info_msg, 0, 0, 1, 2)
|
|
gridlayout.addWidget(train_lbl, 1, 0)
|
|
gridlayout.addWidget(val_lbl, 2, 0)
|
|
gridlayout.addWidget(test_lbl, 3, 0)
|
|
gridlayout.addWidget(self.train_input, 1, 1)
|
|
gridlayout.addWidget(self.val_input, 2, 1)
|
|
gridlayout.addWidget(self.test_input, 3, 1)
|
|
|
|
bb = BB(BB.Ok | BB.Cancel, Qt.Horizontal, self)
|
|
bb.button(BB.Ok).setIcon(newIcon('done'))
|
|
bb.button(BB.Cancel).setIcon(newIcon('undo'))
|
|
bb.accepted.connect(self.validate)
|
|
bb.rejected.connect(self.cancel)
|
|
gridlayout.addWidget(bb, 4, 0, 1, 2)
|
|
|
|
self.setLayout(gridlayout)
|
|
|
|
self.show()
|
|
|
|
def validate(self):
|
|
self.flag_accept = True
|
|
self.accept()
|
|
|
|
def cancel(self):
|
|
self.flag_accept = False
|
|
self.reject()
|
|
|
|
def getStatus(self):
|
|
return self.flag_accept
|
|
|
|
def getDataPartition(self):
|
|
self.train_ratio = int(self.train_input.text())
|
|
self.val_ratio = int(self.val_input.text())
|
|
self.test_ratio = int(self.test_input.text())
|
|
|
|
return self.train_ratio, self.val_ratio, self.test_ratio
|
|
|
|
def closeEvent(self, event):
|
|
self.flag_accept = False
|
|
self.reject()
|
|
|
|
|