diff --git a/config.py b/config.py index 7e10789..2bfe16b 100755 --- a/config.py +++ b/config.py @@ -24,6 +24,7 @@ def get_args(): parser.add_argument('--decensor_input_path', default='./decensor_input/', help='input images with censored regions colored green to be decensored by decensor.py path') parser.add_argument('--decensor_input_original_path', default='./decensor_input_original/', help='input images with no modifications to be decensored by decensor.py path') parser.add_argument('--decensor_output_path', default='./decensor_output/', help='output images generated from running decensor.py path') + parser.add_argument('--clean-up-input-dirs', dest='clean_up_input_dirs', action='store_true', default=False, help='whether to delete all image files in the input directories when decensoring is done') #Decensor settings parser.add_argument('--mask_color_red', default=0, help='red channel of mask color in decensoring') diff --git a/decensor.py b/decensor.py index 50c1dca..7699f76 100755 --- a/decensor.py +++ b/decensor.py @@ -49,6 +49,7 @@ class Decensor(QtCore.QThread): self.decensor_input_path = args.decensor_input_path self.decensor_input_original_path = args.decensor_input_original_path self.decensor_output_path = args.decensor_output_path + self.clean_up_input_dirs = args.clean_up_input_dirs self.signals = IgnoreAll() # Signals class will be given by progressWindow @@ -76,11 +77,15 @@ class Decensor(QtCore.QThread): elif self.warm_up: print("elif not self.warm_up:") self.decensor_all_images_in_folder() + self.do_post_jobs() def stop(self): # in case of stopping decensor, terminate not to run if self while MainWindow is closed self.terminate() + def set_clean_up_input_dirs(self, value): + self.clean_up_input_dirs = value + def find_mask(self, colored): # self.signals.update_progress_LABEL.emit("find_mask()", "finding mask...") mask = np.ones(colored.shape, np.uint8) @@ -373,6 +378,36 @@ class Decensor(QtCore.QThread): print("Decensored image. Returning it.") return output_img + def clean_input_directories(self): + """Removes .png, .jpg, and .jpeg files from input directories.""" + + allowed_extensions = {".png", ".jpg", ".jpeg"} + + self.signals.insertText_progressCursor.emit("Cleaning {}...".format(self.decensor_input_path)) + + for file_name in os.listdir(self.decensor_input_path): + file_path = os.path.join(self.decensor_input_path, file_name) + if os.path.isfile(file_path) and os.path.splitext(file_name)[1].lower() in allowed_extensions: + os.remove(file_path) + + if self.is_mosaic: + self.signals.insertText_progressCursor.emit( + "Cleaning {}...".format(self.decensor_input_original_path) + ) + for file_name in os.listdir(self.decensor_input_original_path): + file_path = os.path.join(self.decensor_input_original_path, + file_name) + if os.path.isfile(file_path) and os.path.splitext(file_name)[1].lower() in allowed_extensions: + os.remove(file_path) + + self.signals.insertText_progressCursor.emit("Done!") + + def do_post_jobs(self): + if self.clean_up_input_dirs: + self.clean_input_directories() + + + if __name__ == '__main__': decensor = Decensor() decensor.decensor_all_images_in_folder() diff --git a/main.py b/main.py index dc2419d..11697d3 100644 --- a/main.py +++ b/main.py @@ -5,7 +5,7 @@ # The greater the number of variations, the longer decensoring process will be. import sys, time -from PySide2.QtWidgets import QWidget, QHBoxLayout, QVBoxLayout, QGridLayout, QGroupBox, QDesktopWidget, QApplication +from PySide2.QtWidgets import QWidget, QHBoxLayout, QVBoxLayout, QGridLayout, QGroupBox, QDesktopWidget, QApplication, QCheckBox from PySide2.QtWidgets import QAction, qApp, QApplication, QMessageBox, QRadioButton, QPushButton, QTextEdit, QLabel from PySide2.QtWidgets import QSizePolicy,QMainWindow, QStatusBar, QProgressBar from PySide2.QtCore import Qt, QObject @@ -23,8 +23,8 @@ class MainWindow(QWidget): super().__init__() self.signals = Signals() self.initUI() - self.setSignals() self.decensor = Decensor(self) + self.setSignals() self.load_model() def initUI(self): @@ -94,14 +94,19 @@ class MainWindow(QWidget): self.statusBar.addWidget(self.statusLabel, 1) self.statusBar.addWidget(self.progressBar, 2) + self.cleanInputDirectoryCheckbox = QCheckBox( + 'Clean up input directories after decensoring' + ) + #put all groups into grid # addWidget(row, column, rowSpan, columnSpan) grid_layout.addWidget(self.tutorialLabel, 0, 0, 1, 2) grid_layout.addWidget(self.censorTypeGroupBox, 1, 0, 1, 1) grid_layout.addWidget(self.variationsGroupBox, 1, 1, 1, 1) - grid_layout.addWidget(self.decensorButton, 2, 0, 1, 2) - grid_layout.addWidget(self.progressMessage, 3, 0, 1, 2) - grid_layout.addWidget(self.statusBar, 4, 0, 1, 2) + grid_layout.addWidget(self.cleanInputDirectoryCheckbox, 2, 0, 1, 2) + grid_layout.addWidget(self.decensorButton, 3, 0, 1, 2) + grid_layout.addWidget(self.progressMessage, 4, 0, 1, 2) + grid_layout.addWidget(self.statusBar, 5, 0, 1, 2) #window size settings self.resize(900, 600) @@ -128,6 +133,9 @@ class MainWindow(QWidget): self.signals.insertText_progressCursor.connect(self.progressMessage.append) self.signals.clear_progressMessage.connect(self.progressMessage.clear) self.signals.appendText_progressMessage.connect(self.progressMessage.append) + self.signals.update_clean_up_input_dirs_flag.connect( + self.decensor.set_clean_up_input_dirs + ) def decensorClicked(self): self.decensorButton.setEnabled(False) @@ -158,6 +166,9 @@ class MainWindow(QWidget): variations = int(vb.text()) self.decensor.variations = variations + self.signals.update_clean_up_input_dirs_flag.emit( + self.cleanInputDirectoryCheckbox.isChecked() + ) self.decensorButton.setEnabled(False) self.decensor.start() diff --git a/signals.py b/signals.py index 55cbda7..ca9e5d4 100644 --- a/signals.py +++ b/signals.py @@ -40,3 +40,5 @@ class Signals(QtCore.QObject): # str : value to change # direct connect to self.progressMessage.append(str) appendText_progressMessage = QtCore.Signal(str) + + update_clean_up_input_dirs_flag = QtCore.Signal(bool)