diff --git a/decensor.py b/decensor.py index 6609871..c686a8f 100755 --- a/decensor.py +++ b/decensor.py @@ -12,6 +12,8 @@ try: import file from model import InpaintNN from libs.utils import * + # for QThread + from PySide2 import QtCore except ImportError as e: print("Error when importing libraries: ", e) @@ -24,9 +26,14 @@ except ImportError as e: # def write(self, text): # self.textWritten.emit(str(text)) - -class Decensor: +''' +print text later on other label telling status, informations ,... +custom_print -> signals."methodname".emit( ... ) later +changing GUI on other thread(not MainWindow) is not allowed +''' +class Decensor(QtCore.QThread): def __init__(self, text_edit = None, text_cursor = None, ui_mode = None): + super().__init__() args = config.get_args() self.is_mosaic = args.is_mosaic self.variations = args.variations @@ -34,7 +41,9 @@ class Decensor: 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.signals = None # Singals class will be given by progressWindow + if ui_mode is not None: self.ui_mode = ui_mode else: @@ -48,23 +57,32 @@ class Decensor: self.text_cursor = text_cursor self.ui_mode = True + def run(self): + self.decensor_all_images_in_folder() + + def stop(self): + # in case of stopping decensor, terminate not to run if self while MainWindow is closed + self.terminate() + def find_mask(self, colored): + self.signals.update_progress_LABEL.emit("find_mask()", "finding mask...") mask = np.ones(colored.shape, np.uint8) i, j = np.where(np.all(colored[0] == self.mask_color, axis=-1)) mask[0, i, j] = 0 return mask def load_model(self): + self.signals.update_progress_LABEL.emit("load_model()", "loading model...") self.model = InpaintNN(bar_model_name = "./models/bar/Train_775000.meta", bar_checkpoint_name = "./models/bar/", mosaic_model_name = "./models/mosaic/Train_290000.meta", mosaic_checkpoint_name = "./models/mosaic/", is_mosaic=self.is_mosaic) - + def decensor_all_images_in_folder(self): #load model once at beginning and reuse same model self.load_model() - + input_color_dir = self.decensor_input_path file_names = os.listdir(input_color_dir) @@ -72,10 +90,19 @@ class Decensor: output_dir = self.decensor_output_path # Change False to True before release --> file.check_file(input_dir, output_dir, True) + self.signals.update_progress_LABEL.emit("file.check_file()", "checking image files and directory...") file_names, self.files_removed = file.check_file(input_dir, output_dir, False) + self.signals.total_ProgressBar_update_MAX_VALUE.emit("set total progress bar MaxValue : "+str(len(file_names)),len(file_names)) + #convert all images into np arrays and put them in a list - for file_name in file_names: + for n, file_name in enumerate(file_names, start = 1): + self.signals.total_ProgressBar_update_VALUE.emit("decensoring {} / {}".format(n, len(file_names)), n) + # singal progress bar value == masks decensored on image , + # e.g) sample image : 17 + self.signals.singal_ProgressBar_update_VALUE.emit("reset value", 0) # set to 0 for every image at start + self.signals.update_progress_LABEL.emit("for-loop, \"for file_name in file_names:\"","decensoring : "+str(file_name)) + color_file_path = os.path.join(input_color_dir, file_name) color_basename, color_ext = os.path.splitext(file_name) if os.path.isfile(color_file_path) and color_ext.casefold() == ".png": @@ -176,9 +203,11 @@ class Decensor: self.custom_print("No green regions detected! Make sure you're using exactly the right color.") return + self.signals.singal_ProgressBar_update_MAX_VALUE.emit("found {} masked regions".format(len(regions)), len(regions)) output_img_array = ori_array[0].copy() for region_counter, region in enumerate(regions, 1): + self.signals.update_progress_LABEL.emit("for-loop, \"for region_counter, region in enumerate(regions, 1):\"","decensoring censor {}/{}".format(region_counter,len(regions))) bounding_box = expand_bounding(ori, region, expand_factor=1.5) crop_img = ori.crop(bounding_box) # crop_img.show() @@ -258,6 +287,7 @@ class Decensor: bounding_height_index = row + bounding_box[1] if (bounding_width_index, bounding_height_index) in region: output_img_array[bounding_height_index][bounding_width_index] = pred_img_array[i,:,:,:][row][col] + self.signals.singal_ProgressBar_update_VALUE.emit("{} out of {} regions decensored.".format(region_counter, len(regions)), region_counter) self.custom_print("{region_counter} out of {region_count} regions decensored.".format(region_counter=region_counter, region_count=len(regions))) output_img_array = output_img_array * 255.0 @@ -269,6 +299,8 @@ class Decensor: output_img = Image.fromarray(output_img_array.astype('uint8')) output_img = self.apply_variant(output_img, variant_number) + self.signals.update_progress_LABEL.emit("finished", "decensoring finished, saving as file...") + if file_name != None: #save the decensored image base_name, ext = os.path.splitext(file_name) @@ -295,3 +327,4 @@ class Decensor: if __name__ == '__main__': decensor = Decensor() decensor.decensor_all_images_in_folder() + # equivalent to decensor.start() (running as QtThread) diff --git a/main.py b/main.py index 5daa7cc..c7f75de 100644 --- a/main.py +++ b/main.py @@ -5,12 +5,14 @@ # 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, QAction, qApp, QApplication, QMessageBox, QRadioButton, QPushButton, QTextEdit, QLabel, QSizePolicy +from PySide2.QtWidgets import QWidget, QHBoxLayout, QVBoxLayout, QGridLayout, QGroupBox, QDesktopWidget, QApplication, QAction, qApp, QApplication, QMessageBox, QRadioButton, QPushButton, QTextEdit, QLabel, QSizePolicy,QMainWindow from PySide2.QtCore import Qt, QObject from PySide2.QtGui import QFont, QTextCursor from decensor import Decensor +from progressWindow import ProgressWindow + class MainWindow(QWidget): def __init__(self): @@ -113,22 +115,30 @@ class MainWindow(QWidget): variations = int(vb.text()) decensor.variations = variations - decensor.decensor_all_images_in_folder() self.decensorButton.setEnabled(True) + self.hide() + self.progress = ProgressWindow(self, decensor = decensor) + # decensor.decensor_all_images_in_folder() + # def showAbout(self): # QMessageBox.about(self, 'About', "DeepCreamPy v2.2.0 \n Developed by deeppomf") - + # #centers the main window def center(self): - qr = self.frameGeometry() cp = QDesktopWidget().availableGeometry().center() qr.moveCenter(cp) self.move(qr.topLeft()) if __name__ == '__main__': + import os + # you could remove this if statement if there's no error without this + if os.name == 'nt': + import PySide2 + pyqt = os.path.dirname(PySide2.__file__) + QApplication.addLibraryPath(os.path.join(pyqt, "plugins")) app = QApplication(sys.argv) ex = MainWindow() - sys.exit(app.exec_()) \ No newline at end of file + sys.exit(app.exec_()) diff --git a/progressWindow.py b/progressWindow.py new file mode 100644 index 0000000..6df2a9d --- /dev/null +++ b/progressWindow.py @@ -0,0 +1,141 @@ +# window for showing progress while decensoring +# spliting windows by main to check Censor Types, Variations, ect will be better for +# understanding flow of code + +import sys +import PySide2 +from PySide2.QtWidgets import (QApplication, QMainWindow, QPushButton, + QToolTip, QLabel, QProgressBar, QAction, qApp) +# from PyQt5.QtCore import QThread +from signals import Signals +import threading +import time + +class ProgressWindow(QMainWindow): + # debug for setting UI + def __init__(self, MainWindow, decensor, debug = False): + super().__init__() + self.width = 700 + self.height = 500 + self.resize(self.width, self.height) + self.setWindowTitle("DeepCreamPy v2.2.0 Decensoring...") + self.initUI() + + # signal class that could share update progress ui from decensor class (Decensor) + self.setSignals() + self.show() + + if not debug: + print("not debug") + # decensor class initialized with options selected from MainWindow + self.decensor = decensor + self.decensor.signals = self.signals + + # to go back to MainWindow after finshed decensoring + self.mainWindow = MainWindow + + self.runDecensor() + + + def initUI(self): + ''' + Must Todo UI: + 1. add goto decensored file button + 2. Two progress bars + 2-1. total images to decenesor (images in ./decensor_input) + 2-2. current decensoring image's censored area (example marmaid image in DCPv2, 2 / 17) + 3. go back to main button + + Could Do UI: + 1. showing live image decensoring (decensored one by one) + ''' + # progress bar showing images left to be decensored + def setProgressBar(): + bar_X = 50 + bar_Y = 300 + bar_width = 600 + bar_height = 30 + + # images waiting to be decensored + self.total_images_ProgressBar = QProgressBar(self) + # setGeometry(left top x cordinate, left top y cordinate, width, height) + self.total_images_ProgressBar.setGeometry(bar_X, bar_Y, bar_width,bar_height ) + self.total_images_ProgressBar.setMaximum(100) + self.total_images_ProgressBar.setValue(0) + + # showing progress of decensored area + self.singal_image_decensor_ProgressBar = QProgressBar(self) + self.singal_image_decensor_ProgressBar.setGeometry(bar_X, bar_Y+80, bar_width,bar_height ) + self.singal_image_decensor_ProgressBar.setMaximum(100) + self.singal_image_decensor_ProgressBar.setValue(0) + + progress_Label_1 = QLabel(self) + progress_Label_1.move(50, 270) + progress_Label_1.setText("Number of your images") + progress_Label_1.resize(progress_Label_1.sizeHint()) + + progress_Label_2 = QLabel(self) + progress_Label_2.move(50, 300 + 50) + progress_Label_2.setText("Number of image censoring") + progress_Label_2.resize(progress_Label_2.sizeHint()) + + self.progress_status_LABEL = QLabel(self) + self.progress_status_LABEL.move(100, 100) + self.progress_status_LABEL.setText("Decensoring...") + self.progress_status_LABEL.resize(self.progress_status_LABEL.sizeHint()) + + setProgressBar() + + def setSignals(self): + self.signals = Signals() + # set signal variable name same as method name preventing confusion + self.signals.total_ProgressBar_update_MAX_VALUE.connect(self.total_ProgressBar_update_MAX_VALUE) + self.signals.total_ProgressBar_update_VALUE.connect(self.total_ProgressBar_update_VALUE) + self.signals.singal_ProgressBar_update_MAX_VALUE.connect(self.singal_ProgressBar_update_MAX_VALUE) + self.signals.singal_ProgressBar_update_VALUE.connect(self.singal_ProgressBar_update_VALUE) + self.signals.update_progress_LABEL.connect(self.update_progress_LABEL) + + # total_images_to_decensor_ProgressBar + def total_ProgressBar_update_MAX_VALUE(self, msg, max): + # print msg for debugging + print(msg) + self.total_images_ProgressBar.setMaximum(max) + + def total_ProgressBar_update_VALUE(self, msg, val): + # print msg for debugging + print(msg) + self.total_images_ProgressBar.setValue(val) + + def singal_ProgressBar_update_MAX_VALUE(self, msg, max): + # print msg for debugging + print(msg) + self.singal_image_decensor_ProgressBar.setMaximum(max) + + def singal_ProgressBar_update_VALUE(self, msg, val): + # print msg for debugging + print(msg) + self.singal_image_decensor_ProgressBar.setValue(val) + + def update_progress_LABEL(self, msg, status): + print(msg) + self.progress_status_LABEL.setText(status) + self.progress_status_LABEL.resize(self.progress_status_LABEL.sizeHint()) + + def runDecensor(self): + # start decensor in other thread, preventing UI Freezing + print("start run") + self.decensor.start() + +if __name__ == "__main__": + # only use for debuging window + + import os + # you could remove this if statement if there's no error without this + if os.name == 'nt': + import PySide2 + pyqt = os.path.dirname(PySide2.__file__) + QApplication.addLibraryPath(os.path.join(pyqt, "plugins")) + app = QApplication(sys.argv) + ex = ProgressWindow(1, 1, debug = True) + ex.show() + sys.exit( app.exec_() ) diff --git a/signals.py b/signals.py new file mode 100644 index 0000000..3af1965 --- /dev/null +++ b/signals.py @@ -0,0 +1,18 @@ +from PySide2 import QtCore + +# Signals used for sharing status between threads(gui thread <-> decensoring thread) +class Signals(QtCore.QObject): + + # str : tells status (print in cmd for debug) + # int : value to change + # usage example in other class(thread) : + # → self.signals.total_ProgressBar_update_MAX_VALUE.emit("update value :"+str(max), max) + total_ProgressBar_update_MAX_VALUE = QtCore.Signal(str, int) + total_ProgressBar_update_VALUE = QtCore.Signal(str, int) + + singal_ProgressBar_update_MAX_VALUE = QtCore.Signal(str, int) + singal_ProgressBar_update_VALUE = QtCore.Signal(str, int) + + # str : tells status (print in cmd for debug) + # str : String to update label + update_progress_LABEL = QtCore.Signal(str, str)