From ad2459ead434164606924c4283325322242485af Mon Sep 17 00:00:00 2001
From: Splendide Imaginarius
 <119545140+Splendide-Imaginarius@users.noreply.github.com>
Date: Wed, 1 Feb 2023 08:33:58 +0000
Subject: [PATCH] Use ffprobe to detect file type, and nb_read_packets as
 nb_frames fallback

This fixes .mkv inputs, and makes the .flv workaround unnecessary.
---
 inference_realesrgan_video.py | 55 +++++++++++++++++++++++++----------
 1 file changed, 40 insertions(+), 15 deletions(-)

diff --git a/inference_realesrgan_video.py b/inference_realesrgan_video.py
index c3c4d1465..3bbdf2c62 100644
--- a/inference_realesrgan_video.py
+++ b/inference_realesrgan_video.py
@@ -10,6 +10,7 @@
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from basicsr.utils.download_util import load_file_from_url
 from os import path as osp
+from pathlib import Path
 from tqdm import tqdm
 
 from realesrgan import RealESRGANer
@@ -23,16 +24,48 @@
     import ffmpeg
 
 
+def guess_type(file):
+    p = Path(file)
+
+    if not p.exists():
+        return None
+
+    if p.is_dir():
+        return 'folder'
+
+    try:
+        ff_info = get_video_meta_info(file)
+        if ff_info['audio'] is not None:
+            return 'video'
+        if ff_info['nb_frames'] is not None:
+            if ff_info['nb_frames'] > 1:
+                return 'video'
+            if ff_info['nb_frames'] == 1:
+                return 'image'
+    except Exception as e:
+        pass
+
+    mt_guess = mimetypes.guess_type(file)[0]
+    if mt_guess is None:
+        return None
+    if mt_guess.startswith('video'):
+        return 'video'
+    if mt_guess.startswith('image'):
+        return 'image'
+
+    return None
+
+
 def get_video_meta_info(video_path):
     ret = {}
-    probe = ffmpeg.probe(video_path)
+    probe = ffmpeg.probe(video_path, count_packets=None)
     video_streams = [stream for stream in probe['streams'] if stream['codec_type'] == 'video']
     has_audio = any(stream['codec_type'] == 'audio' for stream in probe['streams'])
     ret['width'] = video_streams[0]['width']
     ret['height'] = video_streams[0]['height']
     ret['fps'] = eval(video_streams[0]['avg_frame_rate'])
     ret['audio'] = ffmpeg.input(video_path).audio if has_audio else None
-    ret['nb_frames'] = int(video_streams[0]['nb_frames'])
+    ret['nb_frames'] = int(video_streams[0]['nb_frames']) if 'nb_frames' in video_streams[0] else int(video_streams[0]['nb_read_packets'])
     return ret
 
 
@@ -58,12 +91,12 @@ class Reader:
 
     def __init__(self, args, total_workers=1, worker_idx=0):
         self.args = args
-        input_type = mimetypes.guess_type(args.input)[0]
-        self.input_type = 'folder' if input_type is None else input_type
+        input_type = guess_type(args.input)
+        self.input_type = input_type
         self.paths = []  # for image&folder type
         self.audio = None
         self.input_fps = None
-        if self.input_type.startswith('video'):
+        if self.input_type == 'video':
             video_path = get_sub_video(args, total_workers, worker_idx)
             self.stream_reader = (
                 ffmpeg.input(video_path).output('pipe:', format='rawvideo', pix_fmt='bgr24',
@@ -77,7 +110,7 @@ def __init__(self, args, total_workers=1, worker_idx=0):
             self.nb_frames = meta['nb_frames']
 
         else:
-            if self.input_type.startswith('image'):
+            if self.input_type == 'image':
                 self.paths = [args.input]
             else:
                 paths = sorted(glob.glob(os.path.join(args.input, '*')))
@@ -374,15 +407,7 @@ def main():
     args.input = args.input.rstrip('/').rstrip('\\')
     os.makedirs(args.output, exist_ok=True)
 
-    if mimetypes.guess_type(args.input)[0] is not None and mimetypes.guess_type(args.input)[0].startswith('video'):
-        is_video = True
-    else:
-        is_video = False
-
-    if is_video and args.input.endswith('.flv'):
-        mp4_path = args.input.replace('.flv', '.mp4')
-        os.system(f'ffmpeg -i {args.input} -codec copy {mp4_path}')
-        args.input = mp4_path
+    is_video = guess_type(args.input) == 'video'
 
     if args.extract_frame_first and not is_video:
         args.extract_frame_first = False