From 12e9e084c8ce2852a503b088461c070ae346e8c4 Mon Sep 17 00:00:00 2001
From: K4YT3X <k4t@protonmail.com>
Date: Fri, 8 Feb 2019 14:04:43 -0500
Subject: [PATCH] Added ability to specify temp directories

---
 bin/upscaler.py | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/bin/upscaler.py b/bin/upscaler.py
index c8d2934..c85e775 100755
--- a/bin/upscaler.py
+++ b/bin/upscaler.py
@@ -4,7 +4,7 @@
 Name: Video2X Upscaler
 Author: K4YT3X
 Date Created: December 10, 2018
-Last Modified: December 19, 2018
+Last Modified: February 8, 2019
 
 Licensed under the GNU General Public License Version 3 (GNU GPL v3),
     available at: https://www.gnu.org/licenses/gpl-3.0.txt
@@ -30,7 +30,7 @@ MODELS_AVAILABLE = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo',
 
 class Upscaler:
 
-    def __init__(self, input_video, output_video, method, waifu2x_path, ffmpeg_path, ffmpeg_arguments=[], ffmpeg_hwaccel='gpu', output_width=False, output_height=False, factor=False, model_type='anime_style_art_rgb', threads=3):
+    def __init__(self, input_video, output_video, method, waifu2x_path, ffmpeg_path, ffmpeg_arguments=[], ffmpeg_hwaccel='gpu', output_width=False, output_height=False, factor=False, model_type='anime_style_art_rgb', threads=3, extracted_frames=False, upscaled_frames=False, preserve_frames=False):
         # Mandatory arguments
         self.input_video = input_video
         self.output_video = output_video
@@ -48,8 +48,25 @@ class Upscaler:
         self.threads = threads
 
         # Make temporary directories
-        self.extracted_frames = tempfile.mkdtemp()
-        self.upscaled_frames = tempfile.mkdtemp()
+        self.extracted_frames = extracted_frames
+        if not extracted_frames:
+            self.extracted_frames = tempfile.mkdtemp()
+        Avalon.debug_info('Extracted frames is being saved to: {}'.format(self.extracted_frames))
+
+        self.upscaled_frames = upscaled_frames
+        if not upscaled_frames:
+            self.upscaled_frames = tempfile.mkdtemp()
+        Avalon.debug_info('Upscaled frames is being saved to: {}'.format(self.upscaled_frames))
+
+        self.preserve_frames = preserve_frames
+
+    def __del__(self):
+        # Delete temp directories when done
+        if not self.preserve_frames:
+            Avalon.debug_info('Deleting cache directory: {}'.format(self.extracted_frames))
+            shutil.rmtree(self.extracted_frames)
+            Avalon.debug_info('Deleting cache directory: {}'.format(self.upscaled_frames))
+            shutil.rmtree(self.upscaled_frames)
 
     def _get_video_info(self):
         """Gets input video information