| | import cv2 |
| | import gc |
| | import requests |
| | from io import BytesIO |
| | import base64 |
| | from scipy import misc |
| | from PIL import Image |
| | from matplotlib.axes import Axes |
| | from matplotlib.figure import Figure |
| | from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas |
| | from typing import Tuple |
| |
|
| | import torch |
| | from fastai.core import * |
| | from fastai.vision import * |
| |
|
| | from .filters import IFilter, MasterFilter, ColorizerFilter |
| | from .generators import gen_inference_deep, gen_inference_wide |
| |
|
| |
|
| |
|
| | |
| | class ModelImageVisualizer: |
| | def __init__(self, filter: IFilter, results_dir: str = None): |
| | self.filter = filter |
| | self.results_dir = None if results_dir is None else Path(results_dir) |
| | self.results_dir.mkdir(parents=True, exist_ok=True) |
| |
|
| | def _clean_mem(self): |
| | torch.cuda.empty_cache() |
| | |
| |
|
| | def _open_pil_image(self, path: Path) -> Image: |
| | return Image.open(path).convert('RGB') |
| |
|
| | def _get_image_from_url(self, url: str) -> Image: |
| | response = requests.get(url, timeout=30, headers={'Accept': '*/*;q=0.8'}) |
| | img = Image.open(BytesIO(response.content)).convert('RGB') |
| | return img |
| |
|
| | def plot_transformed_image_from_url( |
| | self, |
| | url: str, |
| | path: str = 'test_images/image.png', |
| | results_dir:Path = None, |
| | figsize: Tuple[int, int] = (20, 20), |
| | render_factor: int = None, |
| | |
| | display_render_factor: bool = False, |
| | compare: bool = False, |
| | post_process: bool = True, |
| | watermarked: bool = True, |
| | ) -> Path: |
| | img = self._get_image_from_url(url) |
| | img.save(path) |
| | return self.plot_transformed_image( |
| | path=path, |
| | results_dir=results_dir, |
| | figsize=figsize, |
| | render_factor=render_factor, |
| | display_render_factor=display_render_factor, |
| | compare=compare, |
| | post_process = post_process, |
| | watermarked=watermarked, |
| | ) |
| |
|
| | def plot_transformed_image( |
| | self, |
| | path: str, |
| | results_dir:Path = None, |
| | figsize: Tuple[int, int] = (20, 20), |
| | render_factor: int = None, |
| | display_render_factor: bool = False, |
| | compare: bool = False, |
| | post_process: bool = True, |
| | watermarked: bool = True, |
| | ) -> Path: |
| | path = Path(path) |
| | if results_dir is None: |
| | results_dir = Path(self.results_dir) |
| | result = self.get_transformed_image( |
| | path, render_factor, post_process=post_process,watermarked=watermarked |
| | ) |
| | orig = self._open_pil_image(path) |
| | if compare: |
| | self._plot_comparison( |
| | figsize, render_factor, display_render_factor, orig, result |
| | ) |
| | else: |
| | self._plot_solo(figsize, render_factor, display_render_factor, result) |
| |
|
| | orig.close() |
| | result_path = self._save_result_image(path, result, results_dir=results_dir) |
| | result.close() |
| | return result_path |
| |
|
| | def plot_transformed_pil_image( |
| | self, |
| | input_image: Image, |
| | figsize: Tuple[int, int] = (20, 20), |
| | render_factor: int = None, |
| | display_render_factor: bool = False, |
| | compare: bool = False, |
| | post_process: bool = True, |
| | ) -> Image: |
| |
|
| | result = self.get_transformed_pil_image( |
| | input_image, render_factor, post_process=post_process |
| | ) |
| |
|
| | if compare: |
| | self._plot_comparison( |
| | figsize, render_factor, display_render_factor, input_image, result |
| | ) |
| | else: |
| | self._plot_solo(figsize, render_factor, display_render_factor, result) |
| |
|
| | return result |
| |
|
| | def _plot_comparison( |
| | self, |
| | figsize: Tuple[int, int], |
| | render_factor: int, |
| | display_render_factor: bool, |
| | orig: Image, |
| | result: Image, |
| | ): |
| | fig, axes = plt.subplots(1, 2, figsize=figsize) |
| | self._plot_image( |
| | orig, |
| | axes=axes[0], |
| | figsize=figsize, |
| | render_factor=render_factor, |
| | display_render_factor=False, |
| | ) |
| | self._plot_image( |
| | result, |
| | axes=axes[1], |
| | figsize=figsize, |
| | render_factor=render_factor, |
| | display_render_factor=display_render_factor, |
| | ) |
| |
|
| | def _plot_solo( |
| | self, |
| | figsize: Tuple[int, int], |
| | render_factor: int, |
| | display_render_factor: bool, |
| | result: Image, |
| | ): |
| | fig, axes = plt.subplots(1, 1, figsize=figsize) |
| | self._plot_image( |
| | result, |
| | axes=axes, |
| | figsize=figsize, |
| | render_factor=render_factor, |
| | display_render_factor=display_render_factor, |
| | ) |
| |
|
| | def _save_result_image(self, source_path: Path, image: Image, results_dir = None) -> Path: |
| | if results_dir is None: |
| | results_dir = Path(self.results_dir) |
| | result_path = results_dir / source_path.name |
| | image.save(result_path) |
| | return result_path |
| |
|
| | def get_transformed_image( |
| | self, path: Path, render_factor: int = None, post_process: bool = True, |
| | 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,post_process=post_process |
| | ) |
| |
|
| | return filtered_image |
| |
|
| | def get_transformed_pil_image( |
| | self, input_image: Image, render_factor: int = None, post_process: bool = True, |
| | ) -> Image: |
| | self._clean_mem() |
| | filtered_image = self.filter.filter( |
| | input_image, input_image, render_factor=render_factor,post_process=post_process |
| | ) |
| |
|
| | return filtered_image |
| |
|
| | def _plot_image( |
| | self, |
| | image: Image, |
| | render_factor: int, |
| | axes: Axes = None, |
| | figsize=(20, 20), |
| | display_render_factor = False, |
| | ): |
| | if axes is None: |
| | _, axes = plt.subplots(figsize=figsize) |
| | axes.imshow(np.asarray(image) / 255) |
| | axes.axis('off') |
| | if render_factor is not None and display_render_factor: |
| | plt.text( |
| | 10, |
| | 10, |
| | 'render_factor: ' + str(render_factor), |
| | color='white', |
| | backgroundcolor='black', |
| | ) |
| |
|
| | def _get_num_rows_columns(self, num_images: int, max_columns: int) -> Tuple[int, int]: |
| | columns = min(num_images, max_columns) |
| | rows = num_images // columns |
| | rows = rows if rows * columns == num_images else rows + 1 |
| | return rows, columns |
| |
|
| |
|
| | def get_image_colorizer( |
| | root_folder: Path = Path('./'), render_factor: int = 35, artistic: bool = True |
| | ) -> ModelImageVisualizer: |
| | if artistic: |
| | return get_artistic_image_colorizer(root_folder=root_folder, render_factor=render_factor) |
| | else: |
| | return get_stable_image_colorizer(root_folder=root_folder, render_factor=render_factor) |
| |
|
| |
|
| | def get_stable_image_colorizer( |
| | root_folder: Path = Path('./'), |
| | weights_name: str = 'ColorizeStable_gen', |
| | results_dir='output', |
| | render_factor: int = 35 |
| | ) -> ModelImageVisualizer: |
| | learn = gen_inference_wide(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 vis |
| |
|
| |
|
| | def get_artistic_image_colorizer( |
| | root_folder: Path = Path('./'), |
| | weights_name: str = 'ColorizeArtistic_gen', |
| | results_dir='output', |
| | render_factor: int = 35 |
| | ) -> ModelImageVisualizer: |
| | 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 vis |
