import numpy as np from PIL import Image import os from copy import deepcopy import config from libs.pconv_hybrid_model import PConvUnet from libs.flood_fill import find_regions, expand_bounding class Decensor(): def __init__(self): self.args = config.get_args() self.decensor_mosaic = self.args.is_mosaic self.mask_color = [self.args.mask_color_red/255.0, self.args.mask_color_green/255.0, self.args.mask_color_blue/255.0] if not os.path.exists(self.args.decensor_output_path): os.makedirs(self.args.decensor_output_path) self.load_model() def get_mask(self,ori, width, height): mask = np.zeros(ori.shape, np.uint8) #count = 0 #TODO: change to iterate over all images in batch when implementing batches for row in range(height): for col in range(width): if np.array_equal(ori[0][row][col], self.mask_color): mask[0, row, col] = 1 return 1-mask def load_model(self): self.model = PConvUnet(weight_filepath='data/logs/') self.model.load( r"./models/model.h5", train_bn=False, lr=0.00005 ) def decensor_all_images_in_folder(self): #load model once at beginning and reuse same model #self.load_model() subdir = self.args.decensor_input_path files = os.listdir(subdir) #convert all images into np arrays and put them in a list for file in files: #print(file) file_path = os.path.join(subdir, file) if os.path.isfile(file_path) and os.path.splitext(file_path)[1] == ".png": print("Decensoring the image {file_path}".format(file_path)) censored_img = Image.open(file_path) self.decensor_image(censored_img, file) #decensors one image at a time #TODO: decensor all cropped parts of the same image in a batch (then i need input for ori an array of those images and make additional changes) def decensor_image(self,ori, file_name): width, height = ori.size #save the alpha channel if the image has an alpha channel has_alpha = False alpha_channel = None 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 = np.asarray(ori) ori_array = np.array(ori_array / 255.0) ori_array = np.expand_dims(ori_array, axis = 0) mask = self.get_mask(ori_array, width, height) regions = find_regions(ori) print("Found {region_count} censored regions in this image!".format(region_count = len(regions))) if len(regions) == 0 and not self.decensor_mosaic: print("No green colored regions detected!") return output_img_array = ori_array[0].copy() for region_counter, region in enumerate(regions, 1): bounding_box = expand_bounding(ori, region) 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((512, 512)) crop_img_array = np.asarray(crop_img) crop_img_array = crop_img_array / 255.0 crop_img_array = np.expand_dims(crop_img_array, axis = 0) #resize the mask images mask_img = mask_img.crop(bounding_box) mask_img = mask_img.resize((512, 512)) #mask_img.show() #convert mask_img back to array mask_array = np.asarray(mask_img) mask_array = np.array(mask_array / 255.0) #the mask has been upscaled so there will be values not equal to 0 or 1 #mask_array[mask_array < 0.01] = 0 mask_array[mask_array > 0] = 1 mask_array = np.expand_dims(mask_array, axis = 0) # Run predictions for this batch of images pred_img_array = self.model.predict([crop_img_array, mask_array, mask_array]) pred_img_array = pred_img_array * 255.0 pred_img_array = np.squeeze(pred_img_array, axis = 0) #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_array = np.asarray(pred_img) pred_img_array = pred_img_array / 255.0 # print(pred_img_array.shape) pred_img_array = np.expand_dims(pred_img_array, axis = 0) for i in range(len(ori_array)): if self.decensor_mosaic: output_img_array = pred_img[i] else: for col in range(bounding_width): for row in range(bounding_height): bounding_width = col + bounding_box[0] bounding_height = row + bounding_box[1] if (bounding_width, bounding_height) in region: output_img_array[bounding_height][bounding_width] = pred_img_array[i,:,:,:][row][col] 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 has_alpha: print(output_img_array.shape) print(alpha_channel.shape) output_img_array = np.concatenate((output_img_array, alpha_channel), axis = 2) output_img = Image.fromarray(output_img_array.astype('uint8')) #save the decensored image #file_name, _ = os.path.splitext(file_name) save_path = os.path.join(self.args.decensor_output_path, file_name) output_img.save(save_path) print("Decensored image saved to {save_path}!".format(save_path=save_path)) return if __name__ == '__main__': decensor = Decensor() decensor.decensor_all_images_in_folder()