decensor -> QThread, main minor changes

This commit is contained in:
ccppoo 2019-10-21 17:08:03 +09:00
parent 40ebc0c9e2
commit 1d06049dfb
2 changed files with 54 additions and 11 deletions

View File

@ -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
@ -35,6 +42,8 @@ class Decensor:
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,13 +57,22 @@ 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",
@ -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)

16
main.py
View File

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