Experimental per frame target quality for svt-av1

Av1an/encode.py

@@ -92,6 +92,7 @@ def startup(project: Project, chunk_queue: List[Chunk]):
     project.workers = min(project.workers, clips)
     print(f'\rQueue: {clips} Workers: {project.workers} Passes: {project.passes}\n'
           f'Params: {" ".join(project.video_params)}')
+
     counter = Manager().Counter(project.get_frames(), initial)
     project.counter = counter
 
@@ -119,45 +120,40 @@ def encode(chunk: Chunk, project: Project):
     :param project: The cli project
     :return: None
     """
-    try:
-        st_time = time.time()
+    st_time = time.time()
 
-        chunk_frames = chunk.frames
+    chunk_frames = chunk.frames
 
-        log(f'Enc: {chunk.name}, {chunk_frames} fr\n\n')
+    log(f'Enc: {chunk.name}, {chunk_frames} fr\n\n')
 
-        # Target Quality Mode
-        if project.target_quality:
-            if project.target_quality_method == 'per_shot':
-                per_shot_target_quality_routine(project, chunk)
-            if project.target_quality_method == 'per_frame':
-                per_frame_target_quality_routine(project, chunk)
+    # Target Quality Mode
+    if project.target_quality:
+        if project.target_quality_method == 'per_shot':
+            per_shot_target_quality_routine(project, chunk)
+        if project.target_quality_method == 'per_frame':
+            per_frame_target_quality_routine(project, chunk)
 
-        ENCODERS[project.encoder].on_before_chunk(project, chunk)
+    ENCODERS[project.encoder].on_before_chunk(project, chunk)
 
-        # skip first pass if reusing
-        start = 2 if project.reuse_first_pass and project.passes >= 2 else 1
+    # skip first pass if reusing
+    start = 2 if project.reuse_first_pass and project.passes >= 2 else 1
 
-        # Run all passes for this chunk
-        for current_pass in range(start, project.passes + 1):
-            tqdm_bar(project, chunk, project.encoder, project.counter, chunk_frames, project.passes, current_pass)
+    # Run all passes for this chunk
+    for current_pass in range(start, project.passes + 1):
+        tqdm_bar(project, chunk, project.encoder, project.counter, chunk_frames, project.passes, current_pass)
 
-        ENCODERS[project.encoder].on_after_chunk(project, chunk)
+    ENCODERS[project.encoder].on_after_chunk(project, chunk)
 
-        # get the number of encoded frames, if no check assume it worked and encoded same number of frames
-        encoded_frames = chunk_frames if project.no_check else frame_check_output(chunk, chunk_frames)
+    # get the number of encoded frames, if no check assume it worked and encoded same number of frames
+    encoded_frames = chunk_frames if project.no_check else frame_check_output(chunk, chunk_frames)
 
-        # write this chunk as done if it encoded correctly
-        if encoded_frames == chunk_frames:
-            write_progress_file(Path(project.temp / 'done.json'), chunk, encoded_frames)
+    # write this chunk as done if it encoded correctly
+    if encoded_frames == chunk_frames:
+        write_progress_file(Path(project.temp / 'done.json'), chunk, encoded_frames)
 
-        enc_time = round(time.time() - st_time, 2)
-        log(f'Done: {chunk.name} Fr: {encoded_frames}\n'
-            f'Fps: {round(encoded_frames / enc_time, 4)} Time: {enc_time} sec.\n\n')
-
-    except Exception as e:
-        _, _, exc_tb = sys.exc_info()
-        print(f'Error in encoding loop {e}\nAt line {exc_tb.tb_lineno}')
+    enc_time = round(time.time() - st_time, 2)
+    log(f'Done: {chunk.name} Fr: {encoded_frames}\n'
+        f'Fps: {round(encoded_frames / enc_time, 4)} Time: {enc_time} sec.\n\n')
 
 
 def frame_check_output(chunk: Chunk, expected_frames: int) -> int:

Chunks/chunk.py

@@ -34,7 +34,8 @@ class Chunk:
         self.temp: Path = temp
         self.frames: int = frames
         self.output_ext: str = output_ext
-        self.per_shot_target_quality_cq: Optional[int] = None
+        self.per_shot_target_quality_cq = None
+        self.per_frame_target_quality_q_list = None
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -116,3 +117,17 @@ class Chunk:
         chunk = Chunk(temp, d['index'], d['ffmpeg_gen_cmd'], d['output_ext'], d['size'], d['frames'])
         chunk.per_shot_target_quality_cq = d['per_shot_target_quality_cq']
         return chunk
+
+    def make_q_file(self, q_list):
+        qfile = self.fake_input_path.with_name(f'q_file_{self.name}').with_suffix('.txt')
+        with open(qfile, 'w') as fl:
+            text = ''
+
+            for x in q_list:
+                text += str(x) + '\n'
+
+            fl.write(text)
+        return qfile
+
+
+

Encoders/encoder.py

@@ -81,6 +81,9 @@ class Encoder(ABC):
         :return: match object from re.search matching the number of encoded frames"""
         pass
 
+    def mod_command(self, command, chunk):
+        pass
+
     def make_pipes(self, a: Project, c: Chunk, passes: int, current_pass: int, output: str, man_q: int = None):
         """
         Creates a pipe for the given chunk with the given args
@@ -100,6 +103,9 @@ class Encoder(ABC):
         elif c.per_shot_target_quality_cq:
             enc_cmd = self.man_q(enc_cmd, c.per_shot_target_quality_cq)
 
+        elif c.per_frame_target_quality_q_list:
+            enc_cmd = self.mod_command(enc_cmd, c)
+
         ffmpeg_gen_pipe = subprocess.Popen(c.ffmpeg_gen_cmd, stdout=PIPE, stderr=DEVNULL)
         ffmpeg_pipe = subprocess.Popen(filter_cmd, stdin=ffmpeg_gen_pipe.stdout, stdout=PIPE, stderr=STDOUT)
         pipe = subprocess.Popen(enc_cmd, stdin=ffmpeg_pipe.stdout, stdout=PIPE,

Encoders/svtav1.py

@@ -42,6 +42,25 @@ class SvtAv1(Encoder):
             )
         ]
 
+
+    def mod_command(self, command, chunk) -> Command:
+        """Return new command with q_file
+
+        :param command: old command
+        :param q: q list
+        :return: command with new cq value"""
+        adjusted_command = command.copy()
+
+        i = list_index_of_regex(adjusted_command, r"(--qp|-q)")
+
+        qp_file = chunk.make_q_file(chunk.per_frame_target_quality_q_list)
+
+        new = ['--use-q-file','1', '--qpfile', f'{qp_file.as_posix()}']
+        new_cmd = adjusted_command[:i] + new + adjusted_command[i+2:]
+
+        return new_cmd
+
+
     def man_q(self, command: Command, q: int) -> Command:
         """Return command with new cq value
 

TargetQuality/per_frame.py

@@ -7,7 +7,9 @@ from Av1an.commandtypes import CommandPair, Command
 from Av1an.logger import log
 from VMAF import call_vmaf, read_weighted_vmaf, read_json
 from .target_quality import gen_probes_names, make_pipes, vmaf_probe, weighted_search
+from scipy import interpolate
 import pprint
+import numpy as np
 
 def per_frame_target_quality_routine(project: Project, chunk: Chunk):
     """
@@ -18,7 +20,7 @@ def per_frame_target_quality_routine(project: Project, chunk: Chunk):
     :param chunk: the Chunk
     :return: None
     """
-    chunk.per_frame_target_quality_cq = per_frame_target_quality(chunk, project)
+    chunk.per_frame_target_quality_q_list = per_frame_target_quality(chunk, project)
 
 
 def make_q_file(q_list, chunk):
@@ -28,7 +30,6 @@ def make_q_file(q_list, chunk):
 
         for x in q_list:
             text += str(x) + '\n'
-
         fl.write(text)
     return qfile
 
@@ -66,7 +67,7 @@ def per_frame_probe_cmd(chunk: Chunk, q, ffmpeg_pipe, encoder, probing_rate, qp_
 
 
 def per_frame_probe(q_list, q, chunk, project):
-    qfile = make_q_file(q_list, chunk)
+    qfile = chunk.make_q_file(q_list)
     cmd = per_frame_probe_cmd(chunk, q, project.ffmpeg_pipe, project.encoder, 1, qfile)
     pipe = make_pipes(chunk.ffmpeg_gen_cmd, cmd)
     process_pipe(pipe, chunk)
@@ -85,53 +86,64 @@ def add_probes_to_frame_list(frame_list, q_list, vmafs):
     return frame_list
 
 
+
 def per_frame_target_quality(chunk, project):
     frames = chunk.frames
-    pp = pprint.PrettyPrinter(indent=4).pprint
-    # First q value to make probe at
-    middle_point = (project.min_q + project.max_q) // 2
     frame_list = [{'frame_number': x, 'probes': []} for x in range(frames)]
 
-    # Initial q list
-    for i in range(project.probes):
-        if i == 0:
-            q_list = [ middle_point for x in range(frames)]
-        elif i == 1:
-            q_list = gen_border_probes_q(frame_list, project.min_q, project.max_q, project.target_quality)
-        else:
-            q_list = gen_next_q(frame_list, chunk, project)
-
+    for _ in range(project.probes):
+        q_list = gen_next_q(frame_list, chunk, project)
         vmafs = per_frame_probe(q_list, 1, chunk, project)
         frame_list = add_probes_to_frame_list(frame_list, q_list, vmafs)
+        mse = round(get_square_error([x['probes'][-1][1] for x in frame_list] ,project.target_quality), 2)
+        # print(':: MSE:', mse)
 
-        print(get_square_error([x['probes'][-1][1] for x in frame_list] ,project.target_quality))
-    pp(frame_list)
-    exit()
+        if mse < 1.0:
+            return q_list
+
+    return q_list
 
 
 def get_square_error(ls, target):
-
     total = 0
-
     for i in ls:
         dif = i - target
         total += dif ** 2
-
     mse = total / len(ls)
-
     return mse
 
-def gen_border_probes_q(frame_list, min_q, max_q, target):
+
+def gen_next_q(frame_list, chunk, project):
     q_list = []
 
-    for probe in frame_list:
+    probes = len(frame_list[0]['probes'])
 
-        if probe['probes'][0][1] < target:
-            q_list.append(min_q)
-        else:
-            q_list.append(max_q)
+    if probes == 0:
+        return [project.min_q] * len(frame_list)
+    elif probes == 1:
+        return [project.max_q] * len(frame_list)
+    else:
+        for probe in frame_list:
+
+            x = [x[0] for x in probe['probes']]
+            y = [x[1] for x in probe['probes']]
+
+            if probes > 2:
+                if len(x) != len(set(x)):
+                    q_list.append(probe['probes'][-1][0])
+                    continue
+
+            interpolation = 'quadratic' if probes > 2 else 'linear'
+
+            f = interpolate.interp1d(x, y, kind=interpolation)
+            xnew = np.linspace(min(x), max(x), max(x) - min(x))
+            tl = list(zip(xnew, f(xnew)))
+            q = min(tl, key=lambda l: abs(l[1] - project.target_quality))
+
+            q_list.append(int(round(q[0])))
+
+        return q_list
 
-    return q_list
 
 def search(q1, v1, q2, v2, target):
 
@@ -152,18 +164,6 @@ def search(q1, v1, q2, v2, target):
     return new_point
 
 
-
-def gen_next_q(frame_list, chunk, project):
-    q_list = []
-
-
-    for probe in frame_list:
-        p1, p2 = probe['probes'][-2:]
-        q_list.append(search(p1[0],p1[1],p2[0],p2[1], project.target_quality))
-
-    return q_list
-
-
 """
 def frame_types_probe(chunk: Chunk, q, ffmpeg_pipe, encoder, probing_rate, qp_file) -> CommandPair: