Update rife.py

Pkgs/rife-ncnn-vs/rife.py

@@ -8,30 +8,28 @@ import vapoursynth as vs
 core = vs.core
 
 # Vars from command line (via VSPipe)
-input_path = globals()["input"]
-temp_dir_path = globals()["tmpDir"]
-cache_file = globals()["cache"]
-fps_in = globals()["inFps"]
-fps_out = globals()["outFps"]
-fps_out_resampled = globals()["outFpsRes"]
-res_input = globals()["resIn"]
-res_scaled = globals()["resSc"]
-pad = globals()["pad"]
-frames = globals()["frames"] == 'True'
-dedupe = globals()["dedupe"] == 'True'
-allow_redupe = globals()["redupe"] == 'True'
-match_duration = globals()["matchDur"] == 'True'
-sc_sens = float(globals()["sc"])
-loop = globals()["loop"] == 'True'
-factor = globals()["factor"]
-realtime = globals()["rt"] == 'True'
-perf_osd = realtime and globals()["osd"] == 'True'
-show_frame_nums = globals()["debugFrNums"] == 'True'
-show_vars = globals()["debugVars"] == 'True'
-txt_scale = globals()["txtScale"]
-trim = globals()["trim"]
-override_c_matrix = globals()["cMatrix"]
-alpha = globals()["alpha"] == 'True'
+input_path = globals().get("input", "")
+temp_dir_path = globals().get("tmpDir", "")
+cache_file = globals().get("cache", "")
+fps_in = globals().get("inFps", "")
+fps_out = globals().get("outFps", "")
+fps_out_resampled = globals().get("outFpsRes", "")
+res_scaled = globals().get("resSc", "")
+pad = globals().get("pad", "0x0")
+frames = globals().get("frames", "") == 'True'
+dedupe = globals().get("dedupe", "") == 'True'
+allow_redupe = globals().get("redupe", "") == 'True'
+match_duration = globals().get("matchDur", "") == 'True'
+sc_sens = float(globals().get("sc", ""))
+loop = globals().get("loop", "") == 'True'
+factor = globals().get("factor", "")
+realtime = globals().get("rt", "") == 'True'
+perf_osd = realtime and globals().get("osd", "") == 'True'
+show_frame_nums = globals().get("debugFrNums", "") == 'True'
+show_vars = globals().get("debugVars", "") == 'True'
+trim = globals().get("trim", "")
+override_color_matrix = globals().get("cMatrix", "")
+alpha = globals().get("alpha", "") == 'True'
 
 if alpha:
     show_frame_nums = False
@@ -42,13 +40,15 @@ if alpha:
 frames_dir = os.path.join(temp_dir_path, 'frames')
 inframes_json_path = os.path.join(temp_dir_path, 'input.json')
 frames_vs_json_path = os.path.join(temp_dir_path, 'frames.vs.json')
-vfr_json_path = os.path.join(temp_dir_path, 'frameIndexes.json')
+vfr_resample_json_path = os.path.join(temp_dir_path, 'frameIndexes.json')
 infps_num, infps_den = map(int, fps_in.split('/'))
 outfps_num, outfps_den = map(int, fps_out.split('/'))
 outfps_res_num, outfps_res_den = map(int, fps_out_resampled.split('/'))
-res_in_x, res_in_y = map(int, res_input.split('x')) # Input resolution
 res_scaled_x, res_scaled_y = map(int, res_scaled.split('x')) # Scaled resolution
 pad_x, pad_y = map(int, pad.split('x')) # Padding right/bottom
+txt_scale = max(1, min(res_scaled_x // 1000, 4)) # Text scale = scaled width divided by 1000, rounded to int, and clamped to 1-4
+
+frames_produced_total = 0
 
 # Load frames or video
 if frames:
@@ -63,11 +63,9 @@ else:
     if alpha:
         clip = core.std.PropToClip(clip, prop='_Alpha') # Process only alpha channel
 
-
-src_frames = len(clip)  # Amount of source frames
-
-# if show_frame_nums: clip = core.text.FrameNum(clip, alignment=7, scale=txt_scale)  # Input frame counter
-
+width_src = clip.width  # Input resolution width
+height_src = clip.height  # Input resolution height
+framecount_src = len(clip)  # Amount of source frames
 reordered_clip = clip[0]
 
 # Deduplication
@@ -94,7 +92,7 @@ c_matrix = '709'
 
 try:
     m = first_frame_props._Matrix
-    if override_c_matrix: c_matrix = override_c_matrix
+    if override_color_matrix: c_matrix = override_color_matrix
     elif m == 0:  c_matrix = 'rgb'
     elif m == 4:  c_matrix = 'fcc'
     elif m == 5:  c_matrix = '470bg'
@@ -112,9 +110,9 @@ except:
 # Store color range (same as first frame)
 col_range = 'full' if first_frame_props.get('_ColorRange') == 0 else 'limited'
 
-resize = res_scaled and res_scaled != "0x0" and res_scaled != res_input
-res_w = res_scaled_x if resize else res_in_x
-res_h = res_scaled_y if resize else res_in_y
+resize = res_scaled and res_scaled != "0x0" and res_scaled != f"{width_src}x{height_src}"
+res_w = res_scaled_x if resize else width_src
+res_h = res_scaled_y if resize else height_src
 
 # Scene change detection
 if sc_sens > 0.01:
@@ -127,32 +125,27 @@ if colors == "YUV":
 else:
     clip = core.resize.Bicubic(clip=clip, format=vs.RGBS, width=res_w, height=res_h)
 
-info_str = f"Factor: {factor}\nFPS Inp: {fps_in}\nFPS Out: {fps_out}\nFPS Rsp: {fps_out_resampled}\nRes Inp: {res_input}\nRes Scl: {res_scaled}\nColors: {colors} {col_range}\nDe/Redupe: {dedupe}/{dedupe and allow_redupe}\n"
-info_str += f"Loop: {loop}\nScn Detect: {sc_sens}\nMatch Dur: {match_duration}\nVFR: {os.path.isfile(vfr_json_path)}\nTxtScale: {txt_scale}\nTrim: {trim}"
+info_str = f"FPS Inp: {fps_in}\nFPS Out: {fps_out}\nFPS Rsp: {fps_out_resampled}\nRes Inp: {width_src}x{height_src}\nRes Scl: {res_scaled}\nPad: {pad}\nColors: {colors} {col_range}\nDe/Redupe: {dedupe}/{dedupe and allow_redupe}\n"
+info_str += f"Loop: {loop}\nScn Detect: {sc_sens}\nMatch Dur: {match_duration}\nTrim: {trim}"
 
 # Padding to achieve a compatible resolution (some models need a certain modulo)
 if pad_x > 0 or pad_y > 0:
     clip = core.std.AddBorders(clip, right=pad_x, bottom=pad_y)
 
 pre_interp_frames = len(clip)
-frames_processed_total = 0
 
 # RIFE Variables
-r_mdlpath = globals()["mdlPath"]
-r_gpu = int(globals()["gpu"])
-r_threads = int(globals()["gpuThreads"])
-r_uhd = globals()["uhd"] == 'True'
-r_tta = globals()["tta"] == 'True'
-info_str += f"\nGPU: {r_gpu} ({r_threads} thrds)\nUHD: {r_uhd}\nTTA: {r_tta}\nMatrix: {c_matrix}"
+r_mdlpath = globals().get("mdl", "")
+r_gpu = int(globals().get("gpu", "0"))
+r_threads = int(globals().get("gpuThrds", "1"))
+r_uhd = globals().get("uhd", "") == 'True'
+r_tta = globals().get("tta", "") == 'True'
+info_str += f"\nGPU: {r_gpu} ({r_threads} thrds)\nUHD: {r_uhd}\nMatrix: {c_matrix}"
 
 # OSD (input clip)
 def on_frame_in(n, clip):
-    global frames_processed_total, info_str
-    frames_processed_total += 1
     if show_frame_nums:
-        clip = core.text.Text(clip, text=f"{frames_processed_total:06d}", alignment=7, scale=txt_scale)
-    if show_vars:
-        clip = core.text.Text(clip, text=f"\n\n\n{info_str}", alignment=7, scale=txt_scale)
+        clip = core.text.Text(clip, text=f"IN:  {n:06d}", alignment=7, scale=txt_scale)
     return clip
 
 clip = core.std.FrameEval(clip, functools.partial(on_frame_in, clip=clip))
@@ -161,19 +154,19 @@ clip = core.std.FrameEval(clip, functools.partial(on_frame_in, clip=clip))
 r_fac_num, r_fac_den = map(int, factor.split('/'))
 clip = core.rife.RIFE(clip, factor_num=r_fac_num, factor_den=r_fac_den, model_path=r_mdlpath, gpu_id=(None if r_gpu < 0 else r_gpu), gpu_thread=r_threads, tta=r_tta, uhd=r_uhd, sc=sc_sens > 0.01)
 
+frm_count_after_interp = len(clip)
+
 # Reduplication
-def reorder_clip(clip, frames_vs_json_path):
+if dedupe and allow_redupe and not realtime:
     reordered_clip = clip[0]
     with open(frames_vs_json_path, 'r') as json_file:
         frame_list = json.load(json_file)
         for i in frame_list:
-            if i < clip.num_frames:
-                reordered_clip = reordered_clip + clip[i]
-    # Redupe trim and return the result
-    return reordered_clip.std.Trim(1, reordered_clip.num_frames - 1) # Redupe trim
+            reordered_clip = reordered_clip + clip[i]
+    clip = reordered_clip.std.Trim(1, reordered_clip.num_frames - 1) # Redupe trim
 
-if dedupe and allow_redupe and not realtime:
-    clip = reorder_clip(clip, frames_vs_json_path)
+frm_count_after_redupe = len(clip)
+info_str += f"\nBefore/After RD: {frm_count_after_interp}/{frm_count_after_redupe}"
 
 # Set output format & color matrix
 clip = vs.core.resize.Bicubic(clip, format=vs.YUV444P16, matrix_s=c_matrix) if not alpha else vs.core.resize.Bicubic(clip, format=vs.GRAY8, matrix_s=c_matrix)
@@ -184,7 +177,7 @@ if pad_x > 0 or pad_y > 0:
 
 # Factor rounded to int, minus 1
 end_dupe_count = r_fac_num // r_fac_den - 1
-target_count_match = int(globals()["targetMatch"])
+target_count_match = int(globals().get("targetMatch", ""))
 target_count_true = target_count_match - end_dupe_count
 
 if not dedupe:
@@ -196,7 +189,6 @@ if not dedupe:
 # OSD Variables
 frames_produced_prev = 0
 frames_produced_curr = 0
-frames_produced_total = 0
 last_fps_upd_time = time.time() 
 start_time = last_fps_upd_time 
 
@@ -205,7 +197,9 @@ def on_frame_out(n, clip):
     global start_time, frames_produced_total
     frames_produced_total += 1
     if show_frame_nums:
-        clip = core.text.Text(clip, text=f"\n{frames_produced_total:06d}", alignment=7, scale=txt_scale)
+        clip = core.text.Text(clip, text=f"\nOUT: {n:06d}", alignment=7, scale=txt_scale)
+    if show_vars:
+        clip = core.text.Text(clip, text=f"\n\n\n{info_str}", alignment=7, scale=txt_scale)
     if not perf_osd:
         return clip
     fps_avg_time = 2
@@ -230,18 +224,20 @@ clip = core.std.FrameEval(clip, functools.partial(on_frame_out, clip=clip))
 if show_frame_nums:
     factor_str = f"{r_fac_num / r_fac_den:.2f}"
     # clip = core.text.FrameNum(clip, alignment=9, scale=txt_scale)  # Output frame counter
-    clip = core.text.Text(clip, f"Frames: {src_frames}/{pre_interp_frames} -> {len(clip)} [{factor_str}x]", alignment=9, scale=txt_scale)
+    clip = core.text.Text(clip, f"Frames: {framecount_src}/{pre_interp_frames} -> {len(clip)} [{factor_str}x]", alignment=9, scale=txt_scale)
     clip = core.text.Text(clip, f"Target (match): {target_count_match} - Target (true): {target_count_true} - End Dupes: {end_dupe_count}", alignment=3, scale=txt_scale)
 
 # Frames picked to resample VFR video to fps_out_resampled
-if os.path.isfile(vfr_json_path):
-    with open(vfr_json_path) as json_file:
+if os.path.isfile(vfr_resample_json_path):
+    with open(vfr_resample_json_path) as json_file:
         frame_indexes = json.load(json_file)
         clip = core.std.Splice([clip[i] for i in frame_indexes if i < len(clip)])
-        clip = core.std.AssumeFPS(clip, fpsnum=outfps_res_den, fpsden=outfps_res_num)
+
+# if fps_out_resampled != fps_out and outfps_res_num > 0 and outfps_res_den > 0:
+#     clip = core.std.AssumeFPS(clip, fpsnum=outfps_res_num, fpsden=outfps_res_den)
 
 # Loop video indefinitely for realtime mode
 if realtime and loop:
     clip = clip.std.Loop(0)
 
-clip.set_output()
+clip.set_output()