diff --git a/Av1an/encode.py b/Av1an/encode.py index 38db85b..5ed3112 100644 --- a/Av1an/encode.py +++ b/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: diff --git a/Chunks/chunk.py b/Chunks/chunk.py index c2f4ed0..6335478 100644 --- a/Chunks/chunk.py +++ b/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 + + + diff --git a/Encoders/encoder.py b/Encoders/encoder.py index 5e5c671..2644020 100644 --- a/Encoders/encoder.py +++ b/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, diff --git a/Encoders/svtav1.py b/Encoders/svtav1.py index 569db44..4d405ce 100644 --- a/Encoders/svtav1.py +++ b/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 diff --git a/TargetQuality/per_frame.py b/TargetQuality/per_frame.py index 78a23b4..06ca57b 100644 --- a/TargetQuality/per_frame.py +++ b/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: