From a4a498f314c6506eac5f4cbb889b915ca86eeebf Mon Sep 17 00:00:00 2001 From: NevermindNilas Date: Mon, 17 Jun 2024 14:07:09 +0300 Subject: [PATCH] use alive_bar instead of tqdm --- main.py | 38 +++++++++++--------------------------- requirements-linux.txt | 2 +- requirements-windows.txt | 2 +- src/downloadModels.py | 16 ++++++---------- src/ffmpegSettings.py | 1 - src/getFFMPEG.py | 8 +++----- src/updateScript.py | 26 ++++++++++---------------- 7 files changed, 32 insertions(+), 61 deletions(-) diff --git a/main.py b/main.py index c974c9fe..7ac065bf 100644 --- a/main.py +++ b/main.py @@ -25,8 +25,7 @@ import sys import logging - -from tqdm import tqdm +from alive_progress import alive_bar from concurrent.futures import ThreadPoolExecutor from src.argumentsChecker import argumentChecker from src.getVideoMetadata import getVideoMetadata @@ -144,8 +143,6 @@ def processFrame(self, frame): if self.dedup: if self.dedup_process.run(frame): self.dedupCount += 1 - self.pbar.total -= 1 - self.pbar.refresh() return if self.denoise: @@ -160,10 +157,6 @@ def processFrame(self, frame): ) self.writeBuffer.write(frame) - if self.interpolate: - self.pbar.update(self.interpolate_factor) - else: - self.pbar.update(1) except Exception as e: logging.exception(f"Something went wrong while processing the frames, {e}") @@ -171,29 +164,20 @@ def processFrame(self, frame): def process(self): frameCount = 0 self.dedupCount = 0 - while True: - frame = self.readBuffer.read() - if frame is None: - break - self.processFrame(frame) - frameCount += 1 - + increment = 1 if not self.interpolate else self.interpolate_factor + with alive_bar(self.totalFrames * increment, title='Processing', bar='smooth', unit='frames') as bar: + for _ in range(self.totalFrames): + frame = self.readBuffer.read() + self.processFrame(frame) + frameCount += 1 + bar(increment) + logging.info(f"Processed {frameCount} frames") if self.dedupCount > 0: logging.info(f"Deduplicated {self.dedupCount} frames") - + self.writeBuffer.close() - def start(self): - self.pbar = tqdm( - total=self.totalFrames * self.interpolate_factor if self.interpolate else self.totalFrames, - unit="frames", - unit_scale=True, - colour="green", - desc="🎥 Processing Frames", - bar_format="{l_bar}{bar}|{n_fmt}/{total_fmt} [ETA: {remaining}, {rate_fmt}]", - ) - try: ( self.new_width, @@ -237,7 +221,7 @@ def start(self): audio=self.audio, benchmark=self.benchmark, ) - + with ThreadPoolExecutor(max_workers=3) as executor: executor.submit(self.readBuffer.start) executor.submit(self.process) diff --git a/requirements-linux.txt b/requirements-linux.txt index d4a8d62d..63f5b667 100644 --- a/requirements-linux.txt +++ b/requirements-linux.txt @@ -12,7 +12,7 @@ GPUtil spandrel yt-dlp requests -tqdm +alive-progress tensorrt scikit-image scenedetect[opencv] diff --git a/requirements-windows.txt b/requirements-windows.txt index 9cee5bf5..5e9aaf29 100644 --- a/requirements-windows.txt +++ b/requirements-windows.txt @@ -11,7 +11,7 @@ wmi spandrel yt-dlp requests -tqdm +alive-progress tensorrt tensorrt_cu12 scikit-image diff --git a/src/downloadModels.py b/src/downloadModels.py index cec66d87..b9e778aa 100644 --- a/src/downloadModels.py +++ b/src/downloadModels.py @@ -1,7 +1,7 @@ import os import logging import requests -from tqdm import tqdm +from alive_progress import alive_bar dirPath = os.path.dirname(__file__) weightsDir = os.path.join(dirPath, "weights") @@ -363,15 +363,11 @@ def downloadAndLog(model: str, filename: str, download_url: str, folderPath: str total_size_in_bytes = None logging.error(e) - progress_bar = tqdm( - total=total_size_in_bytes, unit="iB", unit_scale=True, colour="green" - ) - - with open(os.path.join(folderPath, filename), "wb") as file: - for data in response.iter_content(chunk_size=1024): - progress_bar.update(len(data)) - file.write(data) - progress_bar.close() + with alive_bar(total_size_in_bytes, title="Progress", bar="smooth") as bar: + with open(os.path.join(folderPath, filename), "wb") as file: + for data in response.iter_content(chunk_size=1024): + file.write(data) + bar(len(data)) if filename.endswith(".zip"): import zipfile diff --git a/src/ffmpegSettings.py b/src/ffmpegSettings.py index 02bb1f0b..665d44df 100644 --- a/src/ffmpegSettings.py +++ b/src/ffmpegSettings.py @@ -2,7 +2,6 @@ import subprocess import numpy as np import os -import sys import shutil import torch import decord diff --git a/src/getFFMPEG.py b/src/getFFMPEG.py index dafbdf71..5cc9062d 100644 --- a/src/getFFMPEG.py +++ b/src/getFFMPEG.py @@ -4,7 +4,7 @@ import zipfile import tarfile import logging -from tqdm import tqdm +from alive_progress import alive_bar from pathlib import Path import platform @@ -53,12 +53,10 @@ def downloadAndExtractFFMPEG(ffmpegPath): response = requests.get(FFMPEGURL, stream=True) totalSizeInBytes = int(response.headers.get("content-length", 0)) - with tqdm( - total=totalSizeInBytes, unit="iB", unit_scale=True, colour="green" - ) as progress_bar, open(ffmpegArchivePath, "wb") as file: + with alive_bar(total=totalSizeInBytes, title="Downloading FFMPEG", bar="smooth") as bar, open(ffmpegArchivePath, "wb") as file: for data in response.iter_content(chunk_size=1024): - progress_bar.update(len(data)) file.write(data) + bar(len(data)) extractFunc(ffmpegArchivePath, ffmpegDir) diff --git a/src/updateScript.py b/src/updateScript.py index a46a1cdd..247f5da4 100644 --- a/src/updateScript.py +++ b/src/updateScript.py @@ -2,10 +2,9 @@ import logging import json import pathlib -from tqdm import tqdm +from alive_progress import alive_bar import subprocess import os -import shutil from .coloredPrints import green @@ -56,21 +55,16 @@ def updateScript(scriptVersion, mainPath): totalSizeInBytes = int( response.headers.get("content-length", 0) ) - blockSize = 1024 - progressBar = tqdm( + + with alive_bar( total=totalSizeInBytes, - unit="iB", - unit_scale=True, - bar_format="{l_bar}%s{bar}%s{r_bar}" - % ("\033[32m", "\033[0m"), - ) - with open(filePath, "wb") as file: - for data in response.iter_content(blockSize): - progressBar.update(len(data)) + title="Downloading", + bar="smooth", + length=30, + ) as bar, open(filePath, "wb") as file: + for data in response.iter_content(1024): file.write(data) - progressBar.close() - if totalSizeInBytes != 0 and progressBar.n != totalSizeInBytes: - print("ERROR, something went wrong") + bar(len(data)) logging.info( "Download complete, extracting files, depending on the size of the update this might take a while" @@ -96,7 +90,7 @@ def updateScript(scriptVersion, mainPath): filePath.unlink() break - + """ TO:DO - WIP, need to fix this answer = input("Do you wish to copy over the models and ffmpeg files to the new directory? (y/n): ")