Merge pull request #11 from nautics889/feature/clean-input-dirs-checkbox

feat: checkbox for cleaning input dirs (#7)

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')

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()

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.cleanInputDirectoryCheckbox, 2, 0, 1, 2)
-		grid_layout.addWidget(self.progressMessage, 3, 0, 1, 2)
+		grid_layout.addWidget(self.decensorButton, 3, 0, 1, 2)
-		grid_layout.addWidget(self.statusBar, 4, 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()

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)