-
Notifications
You must be signed in to change notification settings - Fork 0
/
whispers2t_batch_gui.py
167 lines (133 loc) · 6.96 KB
/
whispers2t_batch_gui.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
import sys
from pathlib import Path
from PySide6.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, QCheckBox, QLabel, QGroupBox, QMessageBox
from PySide6.QtCore import Qt
import torch
from utilities import get_compute_and_platform_info
from whispers2t_batch_transcriber import Worker
from metrics_bar import MetricsBar
from settings import SettingsGroupBox
def is_nvidia_gpu_available():
return torch.cuda.is_available() and "nvidia" in torch.cuda.get_device_name(0).lower()
class MainWindow(QWidget):
def __init__(self):
super().__init__()
self.initUI()
def initUI(self):
self.setWindowTitle("chintellalaw.com - for non-commercial use")
initial_height = 400 if is_nvidia_gpu_available() else 370
self.setGeometry(100, 100, 680, initial_height)
self.setWindowFlags(self.windowFlags() | Qt.WindowStaysOnTopHint)
main_layout = QVBoxLayout()
transcriberGroupBox = QGroupBox("Batch Transcriber (ctranslate2 edition - huggingface edition coming soon)")
transcriberLayout = QVBoxLayout()
self.dirLabel = QLabel("No directory selected")
transcriberLayout.addWidget(self.dirLabel)
self.progressLabel = QLabel("Status: Idle")
transcriberLayout.addWidget(self.progressLabel)
transcriberGroupBox.setLayout(transcriberLayout)
main_layout.addWidget(transcriberGroupBox)
fileExtensionsGroupBox = QGroupBox("File Extensions")
fileExtensionsLayout = QHBoxLayout()
self.file_extension_checkboxes = []
file_extensions = [".aac", ".amr", ".asf", ".avi", ".flac", ".m4a", ".mkv", ".mp3", ".mp4", ".wav", ".wma"]
for extension in file_extensions:
checkbox = QCheckBox(extension)
checkbox.setChecked(True)
self.file_extension_checkboxes.append(checkbox)
fileExtensionsLayout.addWidget(checkbox)
fileExtensionsGroupBox.setLayout(fileExtensionsLayout)
main_layout.addWidget(fileExtensionsGroupBox)
self.settingsGroupBox = SettingsGroupBox(get_compute_and_platform_info, self)
main_layout.addWidget(self.settingsGroupBox)
selectDirLayout = QHBoxLayout()
self.selectDirButton = QPushButton("Select Directory")
self.selectDirButton.clicked.connect(self.selectDirectory)
selectDirLayout.addWidget(self.selectDirButton)
self.recursiveCheckbox = QCheckBox("Process Sub-Folders?")
selectDirLayout.addWidget(self.recursiveCheckbox)
self.processButton = QPushButton("Process")
self.processButton.clicked.connect(self.processFiles)
selectDirLayout.addWidget(self.processButton)
self.stopButton = QPushButton("Stop")
self.stopButton.setEnabled(False)
self.stopButton.setVisible(True)
self.stopButton.clicked.connect(self.stopProcessing)
selectDirLayout.addWidget(self.stopButton)
main_layout.addLayout(selectDirLayout)
self.metricsBar = MetricsBar()
main_layout.addWidget(self.metricsBar)
self.setLayout(main_layout)
def closeEvent(self, event):
self.metricsBar.stop_metrics_collector()
super().closeEvent(event)
def selectDirectory(self):
dirPath = QFileDialog.getExistingDirectory(self, "Select Directory")
if dirPath:
self.directory = dirPath
self.dirLabel.setText(f"Directory: {dirPath}")
self.processButton.setEnabled(True)
def calculate_files_to_process(self):
selected_extensions = [checkbox.text() for checkbox in self.file_extension_checkboxes if checkbox.isChecked()]
patterns = [f'*{ext}' for ext in selected_extensions]
directory_path = Path(self.directory)
total_files = 0
for pattern in patterns:
if self.recursiveCheckbox.isChecked():
total_files += len(list(directory_path.rglob(pattern)))
else:
total_files += len(list(directory_path.glob(pattern)))
return total_files
def processFiles(self):
if hasattr(self, 'directory'):
total_files = self.calculate_files_to_process()
reply = QMessageBox.question(
self,
'Confirm Process',
f"{total_files} files found that match the selected criteria. Click OK to proceed or Cancel to abort.",
QMessageBox.Ok | QMessageBox.Cancel,
QMessageBox.Ok)
if reply == QMessageBox.Ok:
device = self.settingsGroupBox.computeDeviceComboBox.currentText()
size = self.settingsGroupBox.sizeComboBox.currentText()
quantization = self.settingsGroupBox.quantizationComboBox.currentText()
beam_size = self.settingsGroupBox.beamSizeSlider.value()
batch_size = self.settingsGroupBox.batchSizeSlider.value()
output_format = self.settingsGroupBox.formatComboBox.currentText()
task = self.settingsGroupBox.transcribeTranslateComboBox.currentText()
selected_extensions = [checkbox.text() for checkbox in self.file_extension_checkboxes if checkbox.isChecked()]
self.worker = Worker(directory=self.directory,
recursive=self.recursiveCheckbox.isChecked(),
output_format=output_format,
device=device,
size=size,
quantization=quantization,
beam_size=beam_size,
batch_size=batch_size,
task=task,
selected_extensions=selected_extensions)
self.worker.progress.connect(self.updateProgress)
self.worker.finished.connect(self.workerFinished)
self.worker.start()
self.processButton.setEnabled(False)
self.stopButton.setEnabled(True)
else:
self.updateProgress("Processing cancelled by user.")
else:
self.updateProgress("Please select a directory first.")
def stopProcessing(self):
if hasattr(self, 'worker'):
self.worker.request_stop()
self.stopButton.setEnabled(False)
def updateProgress(self, message):
self.progressLabel.setText(f"Status: {message}")
def workerFinished(self, message):
self.processButton.setEnabled(True)
self.stopButton.setEnabled(False)
self.progressLabel.setText(f"Status: Completed | {message}")
if __name__ == "__main__":
app = QApplication(sys.argv)
app.setStyle("Fusion")
mainWindow = MainWindow()
mainWindow.show()
sys.exit(app.exec())