Update inference_video.py

inference_video.py

@@ -10,23 +10,25 @@ import _thread
 import skvideo.io
 from queue import Queue, Empty
 from model.pytorch_msssim import ssim_matlab
+import shutil
+import tempfile
 
 warnings.filterwarnings("ignore")
 
+
 def transferAudio(sourceVideo, targetVideo):
     import shutil
-    import moviepy.editor
-    tempAudioFileName = "./temp/audio.mkv"
+    import tempfile
+    
+    # generate a unique temp directory for this user
+    unique_temp_dir = tempfile.mkdtemp()
+    tempAudioFileName = os.path.join(unique_temp_dir, "audio.mkv")
 
     # split audio from original video file and store in "temp" directory
     if True:
 
         # clear old "temp" directory if it exits
-        if os.path.isdir("temp"):
-            # remove temp directory
-            shutil.rmtree("temp")
-        # create new "temp" directory
-        os.makedirs("temp")
+        # (This logic is replaced by mkdtemp which guarantees a new empty dir, but keeping comment structure)
         # extract audio from video
         os.system('ffmpeg -y -i "{}" -c:a copy -vn {}'.format(sourceVideo, tempAudioFileName))
 
@@ -36,7 +38,7 @@ def transferAudio(sourceVideo, targetVideo):
     os.system('ffmpeg -y -i "{}" -i {} -c copy "{}"'.format(targetNoAudio, tempAudioFileName, targetVideo))
 
     if os.path.getsize(targetVideo) == 0: # if ffmpeg failed to merge the video and audio together try converting the audio to aac
-        tempAudioFileName = "./temp/audio.m4a"
+        tempAudioFileName = os.path.join(unique_temp_dir, "audio.m4a")
         os.system('ffmpeg -y -i "{}" -c:a aac -b:a 160k -vn {}'.format(sourceVideo, tempAudioFileName))
         os.system('ffmpeg -y -i "{}" -i {} -c copy "{}"'.format(targetNoAudio, tempAudioFileName, targetVideo))
         if (os.path.getsize(targetVideo) == 0): # if aac is not supported by selected format
@@ -51,7 +53,8 @@ def transferAudio(sourceVideo, targetVideo):
         os.remove(targetNoAudio)
 
     # remove temp directory
-    shutil.rmtree("temp")
+    shutil.rmtree(unique_temp_dir)
+
 
 parser = argparse.ArgumentParser(description='Interpolation for a pair of images')
 parser.add_argument('--video', dest='video', type=str, default=None)
@@ -137,7 +140,17 @@ else:
         vid_out_name = args.output
     else:
         vid_out_name = '{}_{}X_{}fps.{}'.format(video_path_wo_ext, args.multi, int(np.round(args.fps)), args.ext)
-    vid_out = cv2.VideoWriter(vid_out_name, fourcc, args.fps, (w, h))
+    # OUTPUT QUALITY SETTINGS
+    # -crf 17 is visually lossless (lower is better quality, higher is lower file size).
+    # -preset slow provides better compression efficiency.
+    outputdict = {
+        '-c:v': 'libx264', 
+        '-crf': '17', 
+        '-preset': 'slow',
+        '-pix_fmt': 'yuv420p'
+    }
+    vid_out = skvideo.io.FFmpegWriter(vid_out_name, inputdict={'-r': str(args.fps)}, outputdict=outputdict)
+
 
 def clear_write_buffer(user_args, write_buffer):
     cnt = 0
@@ -149,7 +162,8 @@ def clear_write_buffer(user_args, write_buffer):
             cv2.imwrite('vid_out/{:0>7d}.png'.format(cnt), item[:, :, ::-1])
             cnt += 1
         else:
-            vid_out.write(item[:, :, ::-1])
+            # skvideo expects RGB (item), whereas cv2 expected BGR (item[:, :, ::-1])
+            vid_out.writeFrame(item)
 
 def build_read_buffer(user_args, read_buffer, videogen):
     try:
@@ -278,7 +292,7 @@ while(not write_buffer.empty()):
     time.sleep(0.1)
 pbar.close()
 if not vid_out is None:
-    vid_out.release()
+    vid_out.close() # skvideo uses .close(), not .release()
 
 # move audio to new video file if appropriate
 if args.png == False and fpsNotAssigned == True and not args.video is None:
@@ -287,4 +301,4 @@ if args.png == False and fpsNotAssigned == True and not args.video is None:
     except:
         print("Audio transfer failed. Interpolated video will have no audio")
         targetNoAudio = os.path.splitext(vid_out_name)[0] + "_noaudio" + os.path.splitext(vid_out_name)[1]
-        os.rename(targetNoAudio, vid_out_name)
+        os.rename(targetNoAudio, vid_out_name)