try: # default library import os, logging, sys, config except ImportError as e: print("Error when importing DEFAULT library : ", e) print("\nIf you made script named [\"os.py\", \"logging.py\", \"sys.py\", \"config.py\"] rename it") print("If not, re-install python or check your Python environment variables") try: # local library import file from model import InpaintNN from libs.utils import * # external library import numpy as np from PIL import Image import tensorflow as tf from PySide2 import QtCore # for QThread except ImportError as e: print("\n"+ '='*20 + " ImportError " + "=" * 20 + "\n") if e.__class__.__name__ == "ModuleNotFoundError": print(e) print("Python libraries are missing. You can install all required libraries by running in the command line (terminal)") print("cpu version : pip install -r requirements-cpu.txt") print("gpu version : pip install -r requirements-gpu.txt") else: print("Error when importing libraries: ", e) print("\nIf pip doesn't work, try update through Anaconda") print("install Anaconda : https://www.anaconda.com/distribution/ \n") class Decensor(QtCore.QThread): def __init__(self, parentThread = None, text_edit = None, text_cursor = None, ui_mode = None): super().__init__(parentThread) args = config.get_args() self.is_mosaic = args.is_mosaic self.variations = args.variations self.mask_color = [args.mask_color_red/255.0, args.mask_color_green/255.0, args.mask_color_blue/255.0] 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 # Signals class will be given by progressWindow self.model = None self.warm_up = False # if ui_mode is not None: # self.ui_mode = ui_mode # else: # self.ui_mode = args.ui_mode # # if self.ui_mode: # self.text_edit = text_edit # self.text_cursor = text_cursor # self.ui_mode = True if not os.path.exists(self.decensor_output_path): os.makedirs(self.decensor_output_path) def run(self): if not self.warm_up : print("if self.warm_up :") self.load_model() return elif self.warm_up: print("elif not self.warm_up:") 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.insertText_progressCursor.emit("Loading model ... please wait ...\n") if self.model is None : 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) self.warm_up = True print("load model finished") self.signals.insertText_progressCursor.emit("Loading model finished!\n") self.signals.update_decensorButton_Text.emit("Decensor Your Images") self.signals.update_decensorButton_Enabled.emit(True) def decensor_all_images_in_folder(self): #load model once at beginning and reuse same model if not self.warm_up : # incase of running by source code self.load_model() input_color_dir = self.decensor_input_path file_names = os.listdir(input_color_dir) input_dir = self.decensor_input_path 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...") self.signals.insertText_progressCursor.emit("Checking image files and directory...\n") 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)) ''' print("set total progress bar MaxValue : "+str(len(file_names))) self.signals.update_ProgressBar_MAX_VALUE.emit(len(file_names)) ''' self.signals.insertText_progressCursor.emit("Decensoring {} image files\n".format(len(file_names))) #convert all images into np arrays and put them in a list for n, file_name in enumerate(file_names, start = 1): # self.signals.total_ProgressBar_update_VALUE.emit("Decensoring {} / {}".format(n, len(file_names)), n) ''' self.update_ProgressBar_SET_VALUE.emit(n) print("Decensoring {} / {}".format(n, len(file_names))) ''' self.signals.insertText_progressCursor.emit("Decensoring image file : {}\n".format(file_name)) # signal progress bar value == masks decensored on image , # e.g) sample image : 17 # self.signals.signal_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": print("--------------------------------------------------------------------------") print("Decensoring the image {}\n".format(color_file_path)) try : colored_img = Image.open(color_file_path) except: print("Cannot identify image file (" +str(color_file_path)+")") self.files_removed.append((color_file_path,3)) # incase of abnormal file format change (ex : text.txt -> text.png) continue #if we are doing a mosaic decensor if self.is_mosaic: #get the original file that hasn't been colored ori_dir = self.decensor_input_original_path test_file_names = os.listdir(ori_dir) #since the original image might not be a png, test multiple file formats valid_formats = {".png", ".jpg", ".jpeg"} for test_file_name in test_file_names: test_basename, test_ext = os.path.splitext(test_file_name) if (test_basename == color_basename) and (test_ext.casefold() in valid_formats): ori_file_path = os.path.join(ori_dir, test_file_name) ori_img = Image.open(ori_file_path) # colored_img.show() self.decensor_image_variations(ori_img, colored_img, file_name) break else: #for...else, i.e if the loop finished without encountering break print("Corresponding original, uncolored image not found in {}".format(color_file_path)) print("Check if it exists and is in the PNG or JPG format.") self.signals.insertText_progressCursor.emit("Corresponding original, uncolored image not found in {}\n".format(color_file_path)) self.signals.insertText_progressCursor.emit("Check if it exists and is in the PNG or JPG format.\n") #if we are doing a bar decensor else: self.decensor_image_variations(colored_img, colored_img, file_name) else: print("--------------------------------------------------------------------------") print("Image can't be found: "+str(color_file_path)) self.signals.insertText_progressCursor.emit("Image can't be found: "+str(color_file_path) + "\n") print("--------------------------------------------------------------------------") if self.files_removed is not None: file.error_messages(None, self.files_removed) print("\nDecensoring complete!") #unload model to prevent memory issues # self.signals.update_progress_LABEL.emit("finished", "Decensoring complete! Close this window and reopen DCP to start a new session.") self.signals.insertText_progressCursor.emit("\nDecensoring complete! remove decensored file before decensoring again not to overwrite") self.signals.update_decensorButton_Enabled.emit(True) tf.reset_default_graph() def decensor_image_variations(self, ori, colored, file_name=None): for i in range(self.variations): self.decensor_image_variation(ori, colored, i, file_name) #create different decensors of the same image by flipping the input image def apply_variant(self, image, variant_number): if variant_number == 0: return image elif variant_number == 1: return image.transpose(Image.FLIP_LEFT_RIGHT) elif variant_number == 2: return image.transpose(Image.FLIP_TOP_BOTTOM) else: return image.transpose(Image.FLIP_LEFT_RIGHT).transpose(Image.FLIP_TOP_BOTTOM) #decensors one image at a time #TODO: decensor all cropped parts of the same image in a batch (then i need input for colored an array of those images and make additional changes) def decensor_image_variation(self, ori, colored, variant_number, file_name): ori = self.apply_variant(ori, variant_number) colored = self.apply_variant(colored, variant_number) width, height = ori.size #save the alpha channel if the image has an alpha channel has_alpha = False if (ori.mode == "RGBA"): has_alpha = True alpha_channel = np.asarray(ori)[:,:,3] alpha_channel = np.expand_dims(alpha_channel, axis =-1) ori = ori.convert('RGB') ori_array = image_to_array(ori) ori_array = np.expand_dims(ori_array, axis = 0) if self.is_mosaic: #if mosaic decensor, mask is empty # mask = np.ones(ori_array.shape, np.uint8) # print(mask.shape) colored = colored.convert('RGB') color_array = image_to_array(colored) color_array = np.expand_dims(color_array, axis = 0) mask = self.find_mask(color_array) mask_reshaped = mask[0,:,:,:] * 255.0 mask_img = Image.fromarray(mask_reshaped.astype('uint8')) # mask_img.show() else: mask = self.find_mask(ori_array) #colored image is only used for finding the regions regions = find_regions(colored.convert('RGB'), [v*255 for v in self.mask_color]) print("Found {region_count} censored regions in this image!".format(region_count = len(regions))) self.signals.insertText_progressCursor.emit("Found {region_count} censored regions in this image!".format(region_count = len(regions))) if len(regions) == 0 and not self.is_mosaic: print("No green (0,255,0) regions detected! Make sure you're using exactly the right color.") self.signals.insertText_progressCursor.emit("No green (0,255,0) regions detected! Make sure you're using exactly the right color.\n") return # self.signals.signal_ProgressBar_update_MAX_VALUE.emit("Found {} masked regions".format(len(regions)), len(regions)) print("Found {} masked regions".format(len(regions))) # self.signals.insertText_progressCursor.emit("Found {} masked regions\n".format(len(regions))) self.signals.update_ProgressBar_MAX_VALUE.emit(len(regions)) self.signals.update_ProgressBar_SET_VALUE.emit(0) output_img_array = ori_array[0].copy() for region_counter, region in enumerate(regions, 1): # self.signals.update_progress_LABEL.emit("\"Decensoring regions in image\"","Decensoring censor {}/{}".format(region_counter,len(regions))) self.signals.insertText_progressCursor.emit("Decensoring regions in image, 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() #convert mask back to image mask_reshaped = mask[0,:,:,:] * 255.0 mask_img = Image.fromarray(mask_reshaped.astype('uint8')) #resize the cropped images crop_img = crop_img.resize((256, 256)) crop_img_array = image_to_array(crop_img) #resize the mask images mask_img = mask_img.crop(bounding_box) mask_img = mask_img.resize((256, 256)) # mask_img.show() #convert mask_img back to array mask_array = image_to_array(mask_img) #the mask has been upscaled so there will be values not equal to 0 or 1 # mask_array[mask_array > 0] = 1 # crop_img_array[..., :-1][mask_array==0] = (0,0,0) if not self.is_mosaic: a, b = np.where(np.all(mask_array == 0, axis = -1)) # print(a,b) # print(crop_img_array[a,b]) # print(crop_img_array[a,b,0]) # print(crop_img_array.shape) # print(type(crop_img_array[0,0])) crop_img_array[a,b,:] = 0. # temp = Image.fromarray((crop_img_array * 255.0).astype('uint8')) # temp.show() crop_img_array = np.expand_dims(crop_img_array, axis = 0) mask_array = np.expand_dims(mask_array, axis = 0) # print(np.amax(crop_img_array)) # print(np.amax(mask_array)) # print(np.amax(masked)) # print(np.amin(crop_img_array)) # print(np.amin(mask_array)) # print(np.amin(masked)) # print(mask_array) crop_img_array = crop_img_array * 2.0 - 1 # mask_array = mask_array / 255.0 # Run predictions for this batch of images pred_img_array = self.model.predict(crop_img_array, crop_img_array, mask_array) pred_img_array = np.squeeze(pred_img_array, axis = 0) pred_img_array = (255.0 * ((pred_img_array + 1.0) / 2.0)).astype(np.uint8) #scale prediction image back to original size bounding_width = bounding_box[2]-bounding_box[0] bounding_height = bounding_box[3]-bounding_box[1] #convert np array to image # print(bounding_width,bounding_height) # print(pred_img_array.shape) pred_img = Image.fromarray(pred_img_array.astype('uint8')) # pred_img.show() pred_img = pred_img.resize((bounding_width, bounding_height), resample = Image.BICUBIC) # pred_img.show() pred_img_array = image_to_array(pred_img) # print(pred_img_array.shape) pred_img_array = np.expand_dims(pred_img_array, axis = 0) # copy the decensored regions into the output image for i in range(len(ori_array)): for col in range(bounding_width): for row in range(bounding_height): bounding_width_index = col + bounding_box[0] 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.signal_ProgressBar_update_VALUE.emit("{} out of {} regions decensored.".format(region_counter, len(regions)), region_counter) self.signals.update_ProgressBar_SET_VALUE.emit(region_counter) self.signals.insertText_progressCursor.emit("{} out of {} regions decensored.\n".format(region_counter, len(regions))) 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 #restore the alpha channel if the image had one if has_alpha: output_img_array = np.concatenate((output_img_array, alpha_channel), axis = 2) output_img = Image.fromarray(output_img_array.astype('uint8')) output_img = self.apply_variant(output_img, variant_number) # self.signals.update_progress_LABEL.emit("current image finished", "Decensoring of current image finished. Saving image...") self.signals.insertText_progressCursor.emit("Decensoring of current image finished. Saving image...") print("current image finished") if file_name != None: #save the decensored image base_name, ext = os.path.splitext(file_name) file_name = base_name + " " + str(variant_number) + ext save_path = os.path.join(self.decensor_output_path, file_name) output_img.save(save_path) print("Decensored image saved to {save_path}!".format(save_path=save_path)) self.signals.insertText_progressCursor.emit("Decensored image saved to {save_path}!".format(save_path=save_path)) self.signals.insertText_progressCursor.emit("="*30) else: # Legacy Code piece ↓, used when DCPv1 had ui with Painting print("Decensored image. Returning it.") return output_img # if __name__ == '__main__': # decensor = Decensor() # decensor.decensor_all_images_in_folder() # equivalent to decensor.start() (running as QtThread)