11import gc
22import os
33import time
4- import argparse
54import asyncio
5+ import argparse
66import vapoursynth
7+ from pathlib import Path
78from functools import partial
89from typing import Union , List , Tuple
9- from utils import GetnativeException , plugin_from_identifier , get_attr , get_source_filter , to_float
10+ from getnative . utils import GetnativeException , plugin_from_identifier , get_attr , get_source_filter , to_float
1011try :
1112 import matplotlib as mpl
1213 import matplotlib .pyplot as pyplot
1920Rework by Infi
2021Original Author: kageru https://gist.github.com/kageru/549e059335d6efbae709e567ed081799
2122Thanks: BluBb_mADe, FichteFoll, stux!, Frechdachs, LittlePox
22-
23- Version: 2.0.0
2423"""
2524
2625core = vapoursynth .core
2726core .add_cache = False
2827imwri = getattr (core , "imwri" , getattr (core , "imwrif" , None ))
29- output_dir = os .path .splitext (os .path .basename (__file__ ))[0 ]
3028_modes = ["bilinear" , "bicubic" , "bl-bc" , "all" ]
3129_descale = plugin_from_identifier (core , "tegaf.asi.xe" )
3230if _descale is None :
@@ -52,12 +50,14 @@ def __init__(self, kernel: str, b: Union[float, int] = 0, c: Union[float, int] =
5250 self .c = c
5351 self .taps = taps
5452 self .plugin = _descale_getnative if try_descale_getnative and _descale_getnative is not None else _descale
55- if self .plugin is not None :
56- self .descaler = getattr (self .plugin , f'De{ self .kernel } ' , None )
57- self .upscaler = getattr (core .resize , self .kernel .title ())
53+ if self .plugin is None :
54+ return # no plugin could be the case for lanczos5 and spline64
55+
56+ self .descaler = getattr (self .plugin , f'De{ self .kernel } ' , None )
57+ self .upscaler = getattr (core .resize , self .kernel .title ())
5858
59- self .check_input ()
60- self .check_for_extra_paras ()
59+ self .check_input ()
60+ self .check_for_extra_paras ()
6161
6262 def check_for_extra_paras (self ):
6363 if self .kernel == 'bicubic' :
@@ -113,8 +113,8 @@ def __repr__(self):
113113
114114
115115class GetNative :
116- def __init__ (self , src , scaler , ar , min_h , max_h , frame , img_out , plot_scaling , plot_format , show_plot , no_save ,
117- steps ):
116+ def __init__ (self , src , scaler , ar , min_h , max_h , frame , mask_out , plot_scaling , plot_format , show_plot , no_save ,
117+ steps , output_dir ):
118118 self .plot_format = plot_format
119119 self .plot_scaling = plot_scaling
120120 self .src = src
@@ -123,10 +123,11 @@ def __init__(self, src, scaler, ar, min_h, max_h, frame, img_out, plot_scaling,
123123 self .ar = ar
124124 self .scaler = scaler
125125 self .frame = frame
126- self .img_out = img_out
126+ self .mask_out = mask_out
127127 self .show_plot = show_plot
128128 self .no_save = no_save
129129 self .steps = steps
130+ self .output_dir = output_dir
130131 self .txt_output = ""
131132 self .resolutions = []
132133 self .filename = self .get_filename ()
@@ -140,15 +141,12 @@ async def run(self):
140141 src_luma32 = core .std .Cache (src_luma32 )
141142
142143 # descale each individual frame
143- descaler = self .scaler .descaler
144- upscaler = self .scaler .upscaler
145- clip_list = []
146- for h in range (self .min_h , self .max_h + 1 , self .steps ):
147- clip_list .append (descaler (src_luma32 , self .getw (h ), h ))
144+ clip_list = [self .scaler .descaler (src_luma32 , self .getw (h ), h )
145+ for h in range (self .min_h , self .max_h + 1 , self .steps )]
148146 full_clip = core .std .Splice (clip_list , mismatch = True )
149- full_clip = upscaler (full_clip , self .getw (src .height ), src .height )
147+ full_clip = self . scaler . upscaler (full_clip , self .getw (src .height ), src .height )
150148 if self .ar != src .width / src .height :
151- src_luma32 = upscaler (src_luma32 , self .getw (src .height ), src .height )
149+ src_luma32 = self . scaler . upscaler (src_luma32 , self .getw (src .height ), src .height )
152150 expr_full = core .std .Expr ([src_luma32 * full_clip .num_frames , full_clip ], 'x y - abs dup 0.015 > swap 0 ?' )
153151 full_clip = core .std .CropRel (expr_full , 5 , 5 , 5 , 5 )
154152 full_clip = core .std .PlaneStats (full_clip )
@@ -159,7 +157,7 @@ async def run(self):
159157 vals = []
160158 full_clip_len = len (full_clip )
161159 for frame_index in range (len (full_clip )):
162- print (f"{ frame_index + 1 } /{ full_clip_len } " , end = "\r " )
160+ print (f"\r { frame_index } /{ full_clip_len - 1 } " , end = "" )
163161 fut = asyncio .ensure_future (asyncio .wrap_future (full_clip .get_frame_async (frame_index )))
164162 tasks_pending .add (fut )
165163 futures [fut ] = frame_index
@@ -170,24 +168,29 @@ async def run(self):
170168 tasks_done , _ = await asyncio .wait (tasks_pending )
171169 vals += [(futures .pop (task ), task .result ().props .PlaneStatsAverage ) for task in tasks_done ]
172170 vals = [v for _ , v in sorted (vals )]
173- ratios , vals , txt_output , best_value = self .analyze_results (vals )
174- if not os .path .isdir (output_dir ):
175- os .mkdir (output_dir )
171+ ratios , vals , best_value = self .analyze_results (vals )
172+ print ("\n " ) # move the cursor, so that you not start at the end of the progress bar
176173
177- plot = self .save_plot (vals )
178- if not self .no_save and self .img_out :
179- self .save_images (src_luma32 )
180-
181- txt_output += 'Raw data:\n Resolution\t | Relative Error\t | Relative difference from last\n '
182- txt_output += '\n ' .join ([
174+ self .txt_output += 'Raw data:\n Resolution\t | Relative Error\t | Relative difference from last\n '
175+ self .txt_output += '\n ' .join ([
183176 f'{ i * self .steps + self .min_h :4d} \t \t | { error :.10f} \t \t | { ratios [i ]:.2f} '
184177 for i , error in enumerate (vals )
185178 ])
186- self .txt_output = txt_output # if anyone needs this later
187179
180+ plot , fig = self .save_plot (vals )
188181 if not self .no_save :
189- with open (f"{ output_dir } /{ self .filename } .txt" , "w" ) as stream :
190- stream .writelines (txt_output )
182+ if not os .path .isdir (self .output_dir ):
183+ os .mkdir (self .output_dir )
184+
185+ print (f"Output Path: { self .output_dir } " )
186+ for fmt in self .plot_format .replace (" " , "" ).split (',' ):
187+ fig .savefig (f'{ self .output_dir } /{ self .filename } .{ fmt } ' )
188+
189+ with open (f"{ self .output_dir } /{ self .filename } .txt" , "w" ) as stream :
190+ stream .writelines (self .txt_output )
191+
192+ if self .mask_out :
193+ self .save_images (src_luma32 )
191194
192195 return best_value , plot , self .resolutions
193196
@@ -223,13 +226,13 @@ def analyze_results(self, vals):
223226 f"Native resolution(s) (best guess): "
224227 f"{ 'p, ' .join ([str (r * self .steps + self .min_h ) for r in self .resolutions ])} p"
225228 )
226- txt_output = (
229+ self . txt_output = (
227230 f"Resize Kernel: { self .scaler } \n "
228231 f"{ best_values } \n "
229232 f"Please check the graph manually for more accurate results\n \n "
230233 )
231234
232- return ratios , vals , txt_output , best_values
235+ return ratios , vals , best_values
233236
234237 # Modified from:
235238 # https://github.com/WolframRhodium/muvsfunc/blob/d5b2c499d1b71b7689f086cd992d9fb1ccb0219e/muvsfunc.py#L5807
@@ -242,13 +245,10 @@ def save_plot(self, vals):
242245 dh_sequence = tuple (range (500 , 1001 , self .steps ))
243246 ticks = tuple (dh for i , dh in enumerate (dh_sequence ) if i % (50 // self .steps ) == 0 )
244247 ax .set (xlabel = "Height" , xticks = ticks , ylabel = "Relative error" , title = self .filename , yscale = "log" )
245- if not self .no_save :
246- for fmt in self .plot_format .replace (" " , "" ).split (',' ):
247- fig .savefig (f'{ output_dir } /{ self .filename } .{ fmt } ' )
248248 if self .show_plot :
249249 plot .show ()
250250
251- return plot
251+ return plot , fig
252252
253253 # Original idea by Chibi_goku http://recensubshq.forumfree.it/?t=64839203
254254 # Vapoursynth port by MonoS @github: https://github.com/MonoS/VS-MaskDetail
@@ -262,15 +262,15 @@ def mask_detail(self, clip, final_width, final_height):
262262
263263 def save_images (self , src_luma32 ):
264264 src = src_luma32
265- first_out = imwri .Write (src , 'png' , f'{ output_dir } /{ self .filename } _source%d.png' )
265+ first_out = imwri .Write (src , 'png' , f'{ self . output_dir } /{ self .filename } _source%d.png' )
266266 first_out .get_frame (0 ) # trick vapoursynth into rendering the frame
267267 for r in self .resolutions :
268- r += self .min_h
268+ r = r * self . steps + self .min_h
269269 image = self .mask_detail (src , self .getw (r ), r )
270- mask_out = imwri .Write (image , 'png' , f'{ output_dir } /{ self .filename } _mask_{ r :d} p%d.png' )
270+ mask_out = imwri .Write (image , 'png' , f'{ self . output_dir } /{ self .filename } _mask_{ r :d} p%d.png' )
271271 mask_out .get_frame (0 )
272272 descale_out = self .scaler .descaler (src , self .getw (r ), r )
273- descale_out = imwri .Write (descale_out , 'png' , f'{ output_dir } /{ self .filename } _{ r :d} p%d.png' )
273+ descale_out = imwri .Write (descale_out , 'png' , f'{ self . output_dir } /{ self .filename } _{ r :d} p%d.png' )
274274 descale_out .get_frame (0 )
275275
276276 def get_filename (self ):
@@ -296,7 +296,12 @@ def getnative(args: Union[List, argparse.Namespace], src: vapoursynth.VideoNode,
296296 if type (args ) == list :
297297 args = parser .parse_args (args )
298298
299- if (args .img or args .img_out ) and imwri is None :
299+ output_dir = Path (args .dir ).resolve ()
300+ if not os .access (output_dir , os .W_OK ):
301+ raise PermissionError (f"Missing write permissions: { output_dir } " )
302+ output_dir = output_dir .joinpath ("results" )
303+
304+ if (args .img or args .mask_out ) and imwri is None :
300305 raise GetnativeException ("imwri not found." )
301306
302307 if _descale_getnative is None :
@@ -334,20 +339,19 @@ def getnative(args: Union[List, argparse.Namespace], src: vapoursynth.VideoNode,
334339 print (f"The image height is { src .height } , going higher is stupid! New max_h { src .height } " )
335340 args .max_h = src .height
336341
337- getn = GetNative (src , scaler , args .ar , args .min_h , args .max_h , args .frame , args .img_out , args .plot_scaling ,
338- args .plot_format , args .show_plot , args .no_save , args .steps )
342+ getn = GetNative (src , scaler , args .ar , args .min_h , args .max_h , args .frame , args .mask_out , args .plot_scaling ,
343+ args .plot_format , args .show_plot , args .no_save , args .steps , output_dir )
339344 try :
340345 loop = asyncio .get_event_loop ()
341346 best_value , plot , resolutions = loop .run_until_complete (getn .run ())
342347 except ValueError as err :
343348 raise GetnativeException (f"Error in getnative: { err } " )
344349
345- content = (
350+ gc .collect ()
351+ print (
346352 f"\n { scaler } AR: { args .ar :.2f} Steps: { args .steps } \n "
347- f"{ best_value } \n \n "
353+ f"{ best_value } \n "
348354 )
349- gc .collect ()
350- print (content )
351355
352356 return resolutions , plot
353357
@@ -378,7 +382,7 @@ def _getnative():
378382
379383 for i , scaler in enumerate (mode ):
380384 if scaler is not None and scaler .plugin is None :
381- print (f"No correct descale version found for { scaler } , continuing with next scaler when available." )
385+ print (f"Warning: No correct descale version found for { scaler } , continuing with next scaler when available." )
382386 continue
383387 getnative (args , src , scaler , first_time = True if i == 0 else False )
384388
@@ -392,18 +396,18 @@ def _getnative():
392396parser .add_argument ('--aspect-ratio' , '-ar' , dest = 'ar' , type = to_float , default = 0 , help = 'Force aspect ratio. Only useful for anamorphic input' )
393397parser .add_argument ('--min-height' , '-min' , dest = "min_h" , type = int , default = 500 , help = 'Minimum height to consider' )
394398parser .add_argument ('--max-height' , '-max' , dest = "max_h" , type = int , default = 1000 , help = 'Maximum height to consider' )
395- parser .add_argument ('--generate-images ' , '-img-out ' , dest = 'img_out ' , action = "store_true" , default = False , help = 'Save detail mask as png' )
399+ parser .add_argument ('--output-mask ' , '-mask ' , dest = 'mask_out ' , action = "store_true" , default = False , help = 'Save detail mask as png' )
396400parser .add_argument ('--plot-scaling' , '-ps' , dest = 'plot_scaling' , type = str .lower , default = 'log' , help = 'Scaling of the y axis. Can be "linear" or "log"' )
397401parser .add_argument ('--plot-format' , '-pf' , dest = 'plot_format' , type = str .lower , default = 'svg' , help = 'Format of the output image. Specify multiple formats separated by commas. Can be svg, png, pdf, rgba, jp(e)g, tif(f), and probably more' )
398402parser .add_argument ('--show-plot-gui' , '-pg' , dest = 'show_plot' , action = "store_true" , default = False , help = 'Show an interactive plot gui window.' )
399- parser .add_argument ('--no-save' , '-ns' , dest = 'no_save' , action = "store_true" , default = False , help = 'Do not save files to disk.' )
403+ parser .add_argument ('--no-save' , '-ns' , dest = 'no_save' , action = "store_true" , default = False , help = 'Do not save files to disk. Disables all output arguments! ' )
400404parser .add_argument ('--is-image' , '-img' , dest = 'img' , action = "store_true" , default = False , help = 'Force image input' )
401405parser .add_argument ('--stepping' , '-steps' , dest = 'steps' , type = int , default = 1 , help = 'This changes the way getnative will handle resolutions. Example steps=3 [500p, 503p, 506p ...]' )
402- if __name__ == '__main__' :
406+ parser .add_argument ('--output-dir' , '-dir' , dest = 'dir' , type = str , default = "" , help = 'Sets the path of the output dir where you want all results to be saved. (/results will always be added as last folder)' )
407+ def main ():
403408 parser .add_argument (dest = 'input_file' , type = str , help = 'Absolute or relative path to the input file' )
404409 parser .add_argument ('--use' , '-u' , default = None , help = 'Use specified source filter e.g. (lsmas.LWLibavSource)' )
405410 parser .add_argument ('--mode' , '-m' , dest = 'mode' , type = str , choices = _modes , default = None , help = 'Choose a predefined mode ["bilinear", "bicubic", "bl-bc", "all"]' )
406-
407411 starttime = time .time ()
408412 _getnative ()
409413 print (f'done in { time .time () - starttime :.2f} s' )
0 commit comments