From 805017fea53a476cd55c7cb8da5925adf3b7388f Mon Sep 17 00:00:00 2001 From: Hellow2 Date: Thu, 15 Jun 2023 09:58:48 +0200 Subject: [PATCH] added ffmpeg progressbar --- src/music_kraken/audio/codec.py | 77 +++++++++++++----------------- src/music_kraken/pages/abstract.py | 14 ++++-- src/music_kraken/pages/youtube.py | 22 +++------ 3 files changed, 49 insertions(+), 64 deletions(-) diff --git a/src/music_kraken/audio/codec.py b/src/music_kraken/audio/codec.py index d2b9fd1..43eed9f 100644 --- a/src/music_kraken/audio/codec.py +++ b/src/music_kraken/audio/codec.py @@ -1,47 +1,18 @@ from typing import List, Tuple -import ffmpeg +from tqdm import tqdm +from ffmpeg_progress_yield import FfmpegProgress import subprocess from ..utils.shared import BITRATE, AUDIO_FORMAT, CODEX_LOGGER as LOGGER from ..objects import Target -def remove_intervals(target: Target, interval_list: List[Tuple[float, float]]): +def correct_codec(target: Target, bitrate_kb: int = BITRATE, audio_format: str = AUDIO_FORMAT, interval_list: List[Tuple[float, float]] = None): if not target.exists: LOGGER.warning(f"Target doesn't exist: {target.file_path}") return - - output_target = Target( - path=target._path, - file=str(target._file) + "." + AUDIO_FORMAT - ) - # https://stackoverflow.com/questions/50594412/cut-multiple-parts-of-a-video-with-ffmpeg - aselect_list: List[str] = [] - - start = 0 - next_start = 0 - for end, next_start in interval_list: - aselect_list.append(f"between(t,{start},{end})") - - start = next_start - - aselect_list.append(f"gte(t,{next_start})") - - select = f"aselect='{'+'.join(aselect_list)}',asetpts=N/SR/TB" - - r = subprocess.run(["ffmpeg", "-i", str(target.file_path), "-af", select, str(output_target.file_path)]) - if r.stderr != "": - LOGGER.debug(r.stderr) - - output_target.copy_content(target) - output_target.delete() - - -def correct_codec(target: Target, bitrate_kb: int = BITRATE, audio_format: str = AUDIO_FORMAT): - if not target.exists: - LOGGER.warning(f"Target doesn't exist: {target.file_path}") - return + interval_list = interval_list or [] bitrate_b = int(bitrate_kb / 1024) @@ -49,18 +20,36 @@ def correct_codec(target: Target, bitrate_kb: int = BITRATE, audio_format: str = path=target._path, file=str(target._file) + "." + audio_format ) + + # get the select thingie + # https://stackoverflow.com/questions/50594412/cut-multiple-parts-of-a-video-with-ffmpeg + aselect_list: List[str] = [] + + start = 0 + next_start = 0 + for end, next_start in interval_list: + aselect_list.append(f"between(t,{start},{end})") + start = next_start + aselect_list.append(f"gte(t,{next_start})") + + select = f"aselect='{'+'.join(aselect_list)}',asetpts=N/SR/TB" + + # build the ffmpeg command + ffmpeg_command = [ + "ffmpeg", + "-i", str(target.file_path), + "-af", select, + "-b", str(bitrate_b), + str(output_target.file_path) + ] - stream = ffmpeg.input(target.file_path) - stream = stream.audio - stream = ffmpeg.output( - stream, - str(output_target.file_path), - audio_bitrate=bitrate_b, - format=audio_format - ) - out, err = ffmpeg.run(stream, quiet=True, overwrite_output=True) - if err != "": - LOGGER.debug(err) + # run the ffmpeg command with a progressbar + ff = FfmpegProgress(ffmpeg_command) + with tqdm(total=100, position=1, desc="ffmpeg") as pbar: + for progress in ff.run_command_with_progress(): + pbar.update(progress - pbar.n) + + LOGGER.debug(ff.stderr) output_target.copy_content(target) output_target.delete() diff --git a/src/music_kraken/pages/abstract.py b/src/music_kraken/pages/abstract.py index 641020a..8f8d8dd 100644 --- a/src/music_kraken/pages/abstract.py +++ b/src/music_kraken/pages/abstract.py @@ -1,7 +1,7 @@ import logging import random from copy import copy -from typing import Optional, Union, Type, Dict, Set, List +from typing import Optional, Union, Type, Dict, Set, List, Tuple import threading from queue import Queue @@ -345,15 +345,16 @@ class Page: file=str(random.randint(0, 999999)) ) - r = self.download_song_to_target(source=sources[0], target=temp_target, desc=song.title) + source = sources[0] + r = self.download_song_to_target(source=source, target=temp_target, desc=song.title) if not r.is_fatal_error: - r.merge(self._post_process_targets(song, temp_target)) + r.merge(self._post_process_targets(song, temp_target, source)) return r - def _post_process_targets(self, song: Song, temp_target: Target) -> DownloadResult: - correct_codec(temp_target) + def _post_process_targets(self, song: Song, temp_target: Target, source: Source) -> DownloadResult: + correct_codec(temp_target, interval_list=self.get_skip_intervals(song, source)) self.post_process_hook(song, temp_target) @@ -371,6 +372,9 @@ class Page: return r + def get_skip_intervals(self, song: Song, source: Source) -> List[Tuple[float, float]]: + return [] + def post_process_hook(self, song: Song, temp_target: Target, **kwargs): pass diff --git a/src/music_kraken/pages/youtube.py b/src/music_kraken/pages/youtube.py index b3ce15b..de5cba6 100644 --- a/src/music_kraken/pages/youtube.py +++ b/src/music_kraken/pages/youtube.py @@ -5,6 +5,8 @@ from enum import Enum import sponsorblock from sponsorblock.errors import HTTPException, NotFoundException +from music_kraken.objects import Song, Source + from ..objects import Source, DatabaseObject, Song, Target from .abstract import Page from ..objects import ( @@ -18,7 +20,6 @@ from ..objects import ( FormattedText, ID3Timestamp ) -from ..audio.codec import remove_intervals from ..connection import Connection from ..utils.support_classes import DownloadResult from ..utils.shared import YOUTUBE_LOGGER, INVIDIOUS_INSTANCE, BITRATE, ENABLE_SPONSOR_BLOCK @@ -384,23 +385,15 @@ class YouTube(Page): return DownloadResult(total=1) return DownloadResult(error_message=f"Streaming to the file went wrong: {endpoint}, {str(target.file_path)}") - def post_process_hook(self, song: Song, temp_target: Target, **kwargs): + def get_skip_intervals(self, song: Song, source: Source) -> List[Tuple[float, float]]: if not ENABLE_SPONSOR_BLOCK: - return + return [] - # find the YouTube id - source_list = song.source_collection.get_sources_from_page(self.SOURCE_TYPE) - if len(source_list) <= 0: - self.LOGGER.warning(f"Couldn't find a youtube source in the post_processing_hook for {song.option_string}") - return - - source = source_list[0] parsed = YouTubeUrl(source.url) if parsed.url_type != YouTubeUrlType.VIDEO: self.LOGGER.warning(f"{source.url} is no video url.") - return + return [] - # call the sponsorblock api, and remove the segments from the audio segments = [] try: segments = self.sponsorblock_client.get_skip_segments(parsed.id) @@ -408,6 +401,5 @@ class YouTube(Page): self.LOGGER.debug(f"No sponsor found for the video {parsed.id}.") except HTTPException as e: self.LOGGER.warning(f"{e}") - - if len(segments) <= 0: - remove_intervals(temp_target, [(segment.start, segment.end) for segment in segments]) + + return [(segment.start, segment.end) for segment in segments]