From 4c48af4fa460a7bebaeb6919bed181fe45ad682f Mon Sep 17 00:00:00 2001
From: k4yt3x <k4yt3x@k4yt3x.com>
Date: Sat, 9 May 2020 19:30:24 -0400
Subject: [PATCH] output intermediate video-only file if stream migration fails

---
 src/upscaler.py        | 33 ++++++++++++++++++++++++++++-----
 src/video2x.yaml       |  1 +
 src/wrappers/ffmpeg.py | 13 ++++++++-----
 3 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/src/upscaler.py b/src/upscaler.py
index 594d79c..06fab3a 100755
--- a/src/upscaler.py
+++ b/src/upscaler.py
@@ -522,14 +522,37 @@ class Upscaler:
                     Avalon.info(_('Converting extracted frames into video'))
 
                     # use user defined output size
-                    self.process_pool.append(fm.convert_video(framerate, f'{self.scale_width}x{self.scale_height}', self.upscaled_frames))
+                    self.process_pool.append(fm.assemble_video(framerate,
+                                                               f'{self.scale_width}x{self.scale_height}',
+                                                               self.upscaled_frames))
                     self._wait()
                     Avalon.info(_('Conversion completed'))
 
-                    # migrate audio tracks and subtitles
-                    Avalon.info(_('Migrating audio tracks and subtitles to upscaled video'))
-                    self.process_pool.append(fm.migrate_audio_tracks_subtitles(self.current_input_video, output_video, self.upscaled_frames))
-                    self._wait()
+                    try:
+                        # migrate audio tracks and subtitles
+                        Avalon.info(_('Migrating audio tracks and subtitles to upscaled video'))
+                        self.process_pool.append(fm.migrate_streams(self.current_input_video,
+                                                                    output_video,
+                                                                    self.upscaled_frames))
+                        self._wait()
+
+                    # if failed to copy streams
+                    # use file with only video stream
+                    except subprocess.CalledProcessError:
+                        Avalon.error(_('Failed to migrate streams'))
+                        Avalon.warning(_('Trying to output video without additional streams'))
+
+                        # construct output file path
+                        output_video_path = output_video.parent / f'{output_video.stem}{fm.intermediate_file_name.suffix}'
+
+                        # if output file already exists, cancel
+                        if output_video_path.exists():
+                            Avalon.error(_('Output video file exists, aborting'))
+
+                        # otherwise, rename intermediate file to the output file
+                        else:
+                            Avalon.info(_('Writing intermediate file to: {}').format(output_video_path.absolute()))
+                            (self.upscaled_frames / fm.intermediate_file_name).rename(output_video_path)
 
                     # destroy temp directories
                     self.cleanup_temp_directories()
diff --git a/src/video2x.yaml b/src/video2x.yaml
index a1c37c9..de2d492 100644
--- a/src/video2x.yaml
+++ b/src/video2x.yaml
@@ -97,6 +97,7 @@ anime4kcpp:
   codec: mp4v # Specify the codec for encoding from mp4v(recommended in Windows), dxva(for Windows), avc1(H264, recommended in Linux), vp09(very slow), hevc(not support in Windowds), av01(not support in Windowds) (string [=mp4v])
 ffmpeg:
   ffmpeg_path: '%LOCALAPPDATA%\video2x\ffmpeg-latest-win64-static\bin'
+  intermediate_file_name: 'intermediate.mkv'
   # step 1: extract all frames from original video
   # into temporary directory
   video_to_frames:
diff --git a/src/wrappers/ffmpeg.py b/src/wrappers/ffmpeg.py
index f4b3444..f1bac56 100644
--- a/src/wrappers/ffmpeg.py
+++ b/src/wrappers/ffmpeg.py
@@ -4,7 +4,7 @@
 Name: Video2X FFmpeg Controller
 Author: K4YT3X
 Date Created: Feb 24, 2018
-Last Modified: May 7, 2020
+Last Modified: May 9, 2020
 
 Description: This class handles all FFmpeg related operations.
 """
@@ -32,7 +32,10 @@ class Ffmpeg:
         self.ffmpeg_path = pathlib.Path(self.ffmpeg_settings['ffmpeg_path'])
         self.ffmpeg_binary = self.ffmpeg_path / 'ffmpeg'
         self.ffmpeg_probe_binary = self.ffmpeg_path / 'ffprobe'
+
+        # video metadata
         self.image_format = image_format
+        self.intermediate_file_name = pathlib.Path(self.ffmpeg_settings['intermediate_file_name'])
         self.pixel_format = None
 
     def get_pixel_formats(self):
@@ -133,7 +136,7 @@ class Ffmpeg:
 
         return(self._execute(execute))
 
-    def convert_video(self, framerate, resolution, upscaled_frames):
+    def assemble_video(self, framerate, resolution, upscaled_frames):
         """Converts images into videos
 
         This method converts a set of images into a video
@@ -177,12 +180,12 @@ class Ffmpeg:
 
         # specify output file location
         execute.extend([
-            upscaled_frames / 'no_audio.mp4'
+            upscaled_frames / self.intermediate_file_name
         ])
 
         return(self._execute(execute))
 
-    def migrate_audio_tracks_subtitles(self, input_video, output_video, upscaled_frames):
+    def migrate_streams(self, input_video, output_video, upscaled_frames):
         """ Migrates audio tracks and subtitles from input video to output video
 
         Arguments:
@@ -198,7 +201,7 @@ class Ffmpeg:
 
         execute.extend([
             '-i',
-            upscaled_frames / 'no_audio.mp4',
+            upscaled_frames / self.intermediate_file_name,
             '-i',
             input_video
         ])