mirror of
https://github.com/Deepshift/DeepCreamPy.git
synced 2024-12-01 04:00:21 +00:00
214 lines
9.1 KiB
Python
214 lines
9.1 KiB
Python
#!/usr/bin/env python3
|
|
|
|
try:
|
|
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.utils import *
|
|
except ImportError as err:
|
|
print("Error: ", err)
|
|
print("Could not import modules. Make sure all dependencies are installed.")
|
|
exit(1)
|
|
|
|
class Decensor:
|
|
|
|
def __init__(self):
|
|
self.args = config.get_args()
|
|
self.is_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, colored):
|
|
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.model = PConvUnet()
|
|
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()
|
|
color_dir = self.args.decensor_input_path
|
|
file_names = os.listdir(color_dir)
|
|
|
|
#convert all images into np arrays and put them in a list
|
|
for file_name in file_names:
|
|
color_file_path = os.path.join(color_dir, file_name)
|
|
color_bn, color_ext = os.path.splitext(file_name)
|
|
if os.path.isfile(color_file_path) and color_ext.casefold() == ".png":
|
|
print("--------------------------------------------------------------------------")
|
|
print("Decensoring the image {}".format(color_file_path))
|
|
colored_img = Image.open(color_file_path)
|
|
#if we are doing a mosaic decensor
|
|
if self.is_mosaic:
|
|
#get the original file that hasn't been colored
|
|
ori_dir = self.args.decensor_input_original_path
|
|
#since the original image might not be a png, test multiple file formats
|
|
valid_formats = {".png", ".jpg", ".jpeg"}
|
|
for test_file_name in os.listdir(ori_dir):
|
|
test_bn, test_ext = os.path.splitext(test_file_name)
|
|
if (test_bn == color_bn) 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(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(ori_file_path))
|
|
print("Check if it exists and is in the PNG or JPG format.")
|
|
else:
|
|
self.decensor_image(colored_img, colored_img, file_name)
|
|
elif not file_name.startswith("."): # warn if not a PNG, but ignore .gitkeep
|
|
print("--------------------------------------------------------------------------")
|
|
print("Unsupported file type (not a PNG): {}".format(color_file_path))
|
|
print("--------------------------------------------------------------------------")
|
|
|
|
#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(self, ori, colored, file_name=None):
|
|
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.get_mask(color_array)
|
|
# mask_reshaped = mask[0,:,:,:] * 255.0
|
|
# mask_img = Image.fromarray(mask_reshaped.astype('uint8'))
|
|
# mask_img.show()
|
|
|
|
else:
|
|
mask = self.get_mask(ori_array)
|
|
|
|
#colored image is only used for finding the regions
|
|
regions = find_regions(colored.convert('RGB'))
|
|
print("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 regions detected! Make sure you're using exactly the right color.")
|
|
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 = image_to_array(crop_img)
|
|
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 = 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
|
|
|
|
if self.is_mosaic:
|
|
a, b = np.where(np.all(mask_array == 0, axis = -1))
|
|
print(a, b)
|
|
coords = [coord for coord in zip(a,b) if ((coord[0] + coord[1]) % 2 == 0)]
|
|
a,b = zip(*coords)
|
|
|
|
mask_array[a,b] = 1
|
|
# mask_array = mask_array * 255.0
|
|
# img = Image.fromarray(mask_array.astype('uint8'))
|
|
# img.show()
|
|
# return
|
|
|
|
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 = 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]
|
|
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'))
|
|
|
|
if file_name != None:
|
|
#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
|
|
else:
|
|
print("Decensored image. Returning it.")
|
|
return output_img
|
|
|
|
if __name__ == '__main__':
|
|
decensor = Decensor()
|
|
decensor.decensor_all_images_in_folder()
|