From 2326e4090bf78cbd7eef677a2cd5a5eb2df3ee9b Mon Sep 17 00:00:00 2001 From: Jason Antic Date: Tue, 18 Feb 2020 17:35:05 -0800 Subject: [PATCH] Adding new palette watermark --- deoldify/visualize.py | 78 ++++++++++++++++++++++++++++++----- resource_images/watermark.png | 3 ++ 2 files changed, 71 insertions(+), 10 deletions(-) create mode 100644 resource_images/watermark.png diff --git a/deoldify/visualize.py b/deoldify/visualize.py index 35ea21e..d3b51da 100644 --- a/deoldify/visualize.py +++ b/deoldify/visualize.py @@ -17,6 +17,27 @@ import base64 from IPython import display as ipythondisplay from IPython.display import HTML from IPython.display import Image as ipythonimage +import cv2 + +# adapted from https://www.pyimagesearch.com/2016/04/25/watermarking-images-with-opencv-and-python/ +def get_watermarked(pil_image: Image) -> Image: + image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR) + (h, w) = image.shape[:2] + image = np.dstack([image, np.ones((h, w), dtype="uint8") * 255]) + full_watermark = cv2.imread('./resource_images/watermark.png', cv2.IMREAD_UNCHANGED) + pct = 0.05 + (fwH, fwW) = full_watermark.shape[:2] + wH = int(pct * h) + wW = int((pct * h / fwH) * fwW) + watermark = cv2.resize(full_watermark, (wH, wW), interpolation=cv2.INTER_AREA) + overlay = np.zeros((h, w, 4), dtype="uint8") + overlay[h - wH - 10 : h - 10, 10 : 10 + wW] = watermark + # blend the two images together using transparent overlays + output = image.copy() + cv2.addWeighted(overlay, 0.5, output, 1.0, 0, output) + rgb_image = cv2.cvtColor(output, cv2.COLOR_BGR2RGB) + final_image = Image.fromarray(rgb_image) + return final_image class ModelImageVisualizer: @@ -45,6 +66,7 @@ class ModelImageVisualizer: render_factor: int = None, display_render_factor: bool = False, compare: bool = False, + watermarked: bool = True, ) -> Path: img = self._get_image_from_url(url) img.save(path) @@ -54,6 +76,7 @@ class ModelImageVisualizer: render_factor=render_factor, display_render_factor=display_render_factor, compare=compare, + watermarked=watermarked, ) def plot_transformed_image( @@ -63,9 +86,12 @@ class ModelImageVisualizer: render_factor: int = None, display_render_factor: bool = False, compare: bool = False, + watermarked: bool = True, ) -> Path: path = Path(path) - result = self.get_transformed_image(path, render_factor) + result = self.get_transformed_image( + path, render_factor, watermarked=watermarked + ) orig = self._open_pil_image(path) if compare: self._plot_comparison( @@ -121,12 +147,18 @@ class ModelImageVisualizer: image.save(result_path) return result_path - def get_transformed_image(self, path: Path, render_factor: int = None) -> Image: + def get_transformed_image( + self, path: Path, render_factor: int = None, watermarked: bool = True + ) -> Image: self._clean_mem() orig_image = self._open_pil_image(path) filtered_image = self.filter.filter( orig_image, orig_image, render_factor=render_factor ) + + if watermarked: + return get_watermarked(filtered_image) + return filtered_image def _plot_image( @@ -200,7 +232,9 @@ class VideoColorizer: str(bwframe_path_template), format='image2', vcodec='mjpeg', qscale=0 ).run(capture_stdout=True) - def _colorize_raw_frames(self, source_path: Path, render_factor: int = None): + def _colorize_raw_frames( + self, source_path: Path, render_factor: int = None, watermarked: bool = True + ): colorframes_folder = self.colorframes_root / (source_path.stem) colorframes_folder.mkdir(parents=True, exist_ok=True) self._purge_images(colorframes_folder) @@ -210,7 +244,7 @@ class VideoColorizer: img_path = bwframes_folder / img if os.path.isfile(str(img_path)): color_image = self.vis.get_transformed_image( - str(img_path), render_factor=render_factor + str(img_path), render_factor=render_factor, watermarked=watermarked ) color_image.save(str(colorframes_folder / img)) @@ -265,26 +299,38 @@ class VideoColorizer: return result_path def colorize_from_url( - self, source_url, file_name: str, render_factor: int = None + self, + source_url, + file_name: str, + render_factor: int = None, + watermarked: bool = True, ) -> Path: source_path = self.source_folder / file_name self._download_video_from_url(source_url, source_path) - return self._colorize_from_path(source_path, render_factor=render_factor) + return self._colorize_from_path( + source_path, render_factor=render_factor, watermarked=watermarked + ) def colorize_from_file_name( - self, file_name: str, render_factor: int = None + self, file_name: str, render_factor: int = None, watermarked: bool = True ) -> Path: source_path = self.source_folder / file_name - return self._colorize_from_path(source_path, render_factor=render_factor) + return self._colorize_from_path( + source_path, render_factor=render_factor, watermarked=watermarked + ) - def _colorize_from_path(self, source_path: Path, render_factor: int = None) -> Path: + def _colorize_from_path( + self, source_path: Path, render_factor: int = None, watermarked: bool = True + ) -> Path: if not source_path.exists(): raise Exception( 'Video at path specfied, ' + str(source_path) + ' could not be found.' ) self._extract_raw_frames(source_path) - self._colorize_raw_frames(source_path, render_factor=render_factor) + self._colorize_raw_frames( + source_path, render_factor=render_factor, watermarked=watermarked + ) return self._build_video(source_path) @@ -292,6 +338,18 @@ def get_video_colorizer(render_factor: int = 21) -> VideoColorizer: return get_stable_video_colorizer(render_factor=render_factor) +def get_artistic_video_colorizer( + root_folder: Path = Path('./'), + weights_name: str = 'ColorizeArtistic_gen', + results_dir='result_images', + render_factor: int = 35, +) -> VideoColorizer: + learn = gen_inference_deep(root_folder=root_folder, weights_name=weights_name) + filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor) + vis = ModelImageVisualizer(filtr, results_dir=results_dir) + return VideoColorizer(vis) + + def get_stable_video_colorizer( root_folder: Path = Path('./'), weights_name: str = 'ColorizeVideo_gen', diff --git a/resource_images/watermark.png b/resource_images/watermark.png new file mode 100644 index 0000000..4dbefe1 --- /dev/null +++ b/resource_images/watermark.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:568488613b9e2addbda770324d67feb956d6c8c56c29285717b15ecfd6e77f03 +size 9210