Skip to content
Snippets Groups Projects
Commit 69c384c3 authored by fima's avatar fima :beers:
Browse files

Merge branch 'loading_progress_bar' into 'main'

Progress bar for loading files works on linux

See merge request !105
parents 26c0cb24 aee9f598
No related branches found
No related tags found
1 merge request!105Progress bar for loading files works on linux
...@@ -26,6 +26,7 @@ import qim3d ...@@ -26,6 +26,7 @@ import qim3d
from qim3d.io.logger import log from qim3d.io.logger import log
from qim3d.utils.internal_tools import get_file_size, sizeof, stringify_path from qim3d.utils.internal_tools import get_file_size, sizeof, stringify_path
from qim3d.utils.system import Memory from qim3d.utils.system import Memory
from qim3d.utils import ProgressBar
dask.config.set(scheduler="processes") dask.config.set(scheduler="processes")
...@@ -770,6 +771,7 @@ def load( ...@@ -770,6 +771,7 @@ def load(
dataset_name=None, dataset_name=None,
return_metadata=False, return_metadata=False,
contains=None, contains=None,
progress_bar:bool = True,
force_load: bool = False, force_load: bool = False,
dim_order=(2, 1, 0), dim_order=(2, 1, 0),
**kwargs, **kwargs,
...@@ -796,6 +798,7 @@ def load( ...@@ -796,6 +798,7 @@ def load(
return_metadata (bool, optional): Specifies whether to return metadata or not. Default is False (only for HDF5 and TXRM/TXM/XRM files) return_metadata (bool, optional): Specifies whether to return metadata or not. Default is False (only for HDF5 and TXRM/TXM/XRM files)
contains (str, optional): Specifies a part of the name that is common for the TIFF file stack to be loaded (only for TIFF stacks). contains (str, optional): Specifies a part of the name that is common for the TIFF file stack to be loaded (only for TIFF stacks).
Default is None. Default is None.
progress_bar (bool, optional): Displays tqdm progress bar. Useful for large files. So far works only for linux. Default is False.
force_load (bool, optional): If the file size exceeds available memory, a MemoryError is raised. force_load (bool, optional): If the file size exceeds available memory, a MemoryError is raised.
If force_load is True, the error is changed to warning and the loader tries to load it anyway. Default is False. If force_load is True, the error is changed to warning and the loader tries to load it anyway. Default is False.
dim_order (tuple, optional): The order of the dimensions in the volume for .vol files. Default is (2,1,0) which corresponds to (z,y,x) dim_order (tuple, optional): The order of the dimensions in the volume for .vol files. Default is (2,1,0) which corresponds to (z,y,x)
...@@ -829,6 +832,10 @@ def load( ...@@ -829,6 +832,10 @@ def load(
**kwargs, **kwargs,
) )
if progress_bar and os.name == 'posix':
with ProgressBar(path):
data = loader.load(path)
else:
data = loader.load(path) data = loader.load(path)
def log_memory_info(data): def log_memory_info(data):
......
...@@ -5,4 +5,5 @@ from .data import Dataset, prepare_dataloaders, prepare_datasets ...@@ -5,4 +5,5 @@ from .data import Dataset, prepare_dataloaders, prepare_datasets
from .img import generate_volume, overlay_rgb_images from .img import generate_volume, overlay_rgb_images
from .models import inference, model_summary, train_model from .models import inference, model_summary, train_model
from .preview import image_preview from .preview import image_preview
from .loading_progress_bar import ProgressBar
from .system import Memory from .system import Memory
from threading import Timer
import psutil
import sys
from tqdm.auto import tqdm
from qim3d.utils.internal_tools import get_file_size
class RepeatTimer(Timer):
"""
If the memory check is set as a normal thread, there is no garuantee it will switch
resulting in not enough memory checks to create smooth progress bar or to make it
work at all.
Thus we have to use timer, which runs the function at (approximately) the given time. With this subclass
from https://stackoverflow.com/a/48741004/11514359
we don't have to guess how many timers we will need and create multiple timers.
"""
def run(self):
while not self.finished.wait(self.interval):
self.function(*self.args, **self.kwargs)
class ProgressBar:
def __init__(self, filename:str, repeat_time:float = 0.5, *args, **kwargs):
"""
Creates class for 'with' statement to track progress during loading a file into memory
Parameters:
------------
- filename (str): to get size of the file
- repeat_time (float, optional): How often the timer checks how many bytes were loaded. Even if very small,
it doesn't make the progress bar smoother as there are only few visible changes in number of read_chars.
Defaults to 0.25
"""
self.timer = RepeatTimer(repeat_time, self.memory_check)
self.pbar = tqdm(total = get_file_size(filename),
desc = "Loading: ",
unit = "B",
file = sys.stdout,
unit_scale = True,
unit_divisor = 1024,
bar_format = '{l_bar}{bar}| {n_fmt}{unit}/{total_fmt}{unit} [{elapsed}<{remaining}, ' '{rate_fmt}{postfix}]')
self.last_memory = 0
self.process = psutil.Process()
def memory_check(self):
counters = self.process.io_counters()
try:
memory = counters.read_chars
except AttributeError:
memory = counters.read_bytes + counters.other_bytes
try:
self.pbar.update(memory - self.last_memory)
except AttributeError: # When we leave the context manager, we delete the pbar so it can not be updated anymore
# It's because it takes quite a long time for the timer to end and might update the pbar
# one more time before ending which messes up the whole thing
pass
self.last_memory = memory
def __enter__(self):
self.timer.start()
def __exit__(self, exception_type, exception_value, exception_traceback):
self.timer.cancel()
self.pbar.clear()
self.pbar.n = self.pbar.total
self.pbar.display()
del self.pbar # So the update process can not update it anymore
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment