first commit in hf_spaces

deoldify/visualize.py

@@ -1,51 +1,21 @@
-from fastai.core import *
-from fastai.vision import *
-from matplotlib.axes import Axes
-from matplotlib.figure import Figure
-from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
-from .filters import IFilter, MasterFilter, ColorizerFilter
-from .generators import gen_inference_deep, gen_inference_wide
-# from tensorboardX import SummaryWriter
-from scipy import misc
-from PIL import Image
-# import ffmpeg
-# import youtube_dl
+import cv2
 import gc
 import requests
 from io import BytesIO
 import base64
-# from IPython import display as ipythondisplay
-# from IPython.display import HTML
-# from IPython.display import Image as ipythonimage
-import cv2
+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
+
+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
 
-# # adapted from https://www.pyimagesearch.com/2016/04/25/watermarking-images-with-opencv-and-python/
-# def get_watermarked(pil_image: Image) -> Image:
-#     try:
-#         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])
-#         pct = 0.05
-#         full_watermark = cv2.imread(
-#             './resource_images/watermark.png', cv2.IMREAD_UNCHANGED
-#         )
-#         (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")
-#         (wH, wW) = watermark.shape[:2]
-#         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
-#     except:
-#         # Don't want this to crash everything, so let's just not watermark the image for now.
-#         return pil_image
 
 
 class ModelImageVisualizer:
@@ -245,184 +215,6 @@ class ModelImageVisualizer:
         return rows, columns
 
 
-# class VideoColorizer:
-#     def __init__(self, vis: ModelImageVisualizer):
-#         self.vis = vis
-#         workfolder = Path('./video')
-#         self.source_folder = workfolder / "source"
-#         self.bwframes_root = workfolder / "bwframes"
-#         self.audio_root = workfolder / "audio"
-#         self.colorframes_root = workfolder / "colorframes"
-#         self.result_folder = workfolder / "result"
-
-#     def _purge_images(self, dir):
-#         for f in os.listdir(dir):
-#             if re.search('.*?\.jpg', f):
-#                 os.remove(os.path.join(dir, f))
-
-#     def _get_fps(self, source_path: Path) -> str:
-#         probe = ffmpeg.probe(str(source_path))
-#         stream_data = next(
-#             (stream for stream in probe['streams'] if stream['codec_type'] == 'video'),
-#             None,
-#         )
-#         return stream_data['avg_frame_rate']
-
-#     def _download_video_from_url(self, source_url, source_path: Path):
-#         if source_path.exists():
-#             source_path.unlink()
-
-#         ydl_opts = {
-#             'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/mp4',
-#             'outtmpl': str(source_path),
-#             'retries': 30,
-#             'fragment-retries': 30
-#         }
-#         with youtube_dl.YoutubeDL(ydl_opts) as ydl:
-#             ydl.download([source_url])
-
-#     def _extract_raw_frames(self, source_path: Path):
-#         bwframes_folder = self.bwframes_root / (source_path.stem)
-#         bwframe_path_template = str(bwframes_folder / '%5d.jpg')
-#         bwframes_folder.mkdir(parents=True, exist_ok=True)
-#         self._purge_images(bwframes_folder)
-#         ffmpeg.input(str(source_path)).output(
-#             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, post_process: bool = True,
-#         watermarked: bool = True,
-#     ):
-#         colorframes_folder = self.colorframes_root / (source_path.stem)
-#         colorframes_folder.mkdir(parents=True, exist_ok=True)
-#         self._purge_images(colorframes_folder)
-#         bwframes_folder = self.bwframes_root / (source_path.stem)
-
-#         for img in progress_bar(os.listdir(str(bwframes_folder))):
-#             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, post_process=post_process,watermarked=watermarked
-#                 )
-#                 color_image.save(str(colorframes_folder / img))
-
-#     def _build_video(self, source_path: Path) -> Path:
-#         colorized_path = self.result_folder / (
-#             source_path.name.replace('.mp4', '_no_audio.mp4')
-#         )
-#         colorframes_folder = self.colorframes_root / (source_path.stem)
-#         colorframes_path_template = str(colorframes_folder / '%5d.jpg')
-#         colorized_path.parent.mkdir(parents=True, exist_ok=True)
-#         if colorized_path.exists():
-#             colorized_path.unlink()
-#         fps = self._get_fps(source_path)
-
-#         ffmpeg.input(
-#             str(colorframes_path_template),
-#             format='image2',
-#             vcodec='mjpeg',
-#             framerate=fps,
-#         ).output(str(colorized_path), crf=17, vcodec='libx264').run(capture_stdout=True)
-
-#         result_path = self.result_folder / source_path.name
-#         if result_path.exists():
-#             result_path.unlink()
-#         # making copy of non-audio version in case adding back audio doesn't apply or fails.
-#         shutil.copyfile(str(colorized_path), str(result_path))
-
-#         # adding back sound here
-#         audio_file = Path(str(source_path).replace('.mp4', '.aac'))
-#         if audio_file.exists():
-#             audio_file.unlink()
-
-#         os.system(
-#             'ffmpeg -y -i "'
-#             + str(source_path)
-#             + '" -vn -acodec copy "'
-#             + str(audio_file)
-#             + '"'
-#         )
-
-#         if audio_file.exists:
-#             os.system(
-#                 'ffmpeg -y -i "'
-#                 + str(colorized_path)
-#                 + '" -i "'
-#                 + str(audio_file)
-#                 + '" -shortest -c:v copy -c:a aac -b:a 256k "'
-#                 + str(result_path)
-#                 + '"'
-#             )
-#         print('Video created here: ' + str(result_path))
-#         return result_path
-
-#     def colorize_from_url(
-#         self,
-#         source_url,
-#         file_name: str,
-#         render_factor: int = None,
-#         post_process: bool = True,
-#         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, post_process=post_process,watermarked=watermarked
-#         )
-
-#     def colorize_from_file_name(
-#         self, file_name: str, render_factor: int = None,  watermarked: bool = True, post_process: bool = True,
-#     ) -> Path:
-#         source_path = self.source_folder / file_name
-#         return self._colorize_from_path(
-#             source_path, render_factor=render_factor,  post_process=post_process,watermarked=watermarked
-#         )
-
-#     def _colorize_from_path(
-#         self, source_path: Path, render_factor: int = None,  watermarked: bool = True, post_process: 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,post_process=post_process,watermarked=watermarked
-#         )
-#         return self._build_video(source_path)
-
-
-# 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',
-#     results_dir='result_images',
-#     render_factor: int = 21
-# ) -> VideoColorizer:
-#     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 VideoColorizer(vis)
-
-
 def get_image_colorizer(
     root_folder: Path = Path('./'), render_factor: int = 35, artistic: bool = True
 ) -> ModelImageVisualizer: