Skip to content

Commit

Permalink
Adjust tqdm for batch
Browse files Browse the repository at this point in the history
  • Loading branch information
JasonGross committed Nov 8, 2024
1 parent 73373ef commit e79f40a
Showing 1 changed file with 23 additions and 2 deletions.
25 changes: 23 additions & 2 deletions gbmi/utils/images.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,8 +91,12 @@ def batch_run(
stdout_write: Optional[Callable] = None,
stderr_write: Optional[Callable] = None,
wrap_errs: Optional[list[Exception]] = None,
tqdm_position: Optional[int] = 0,
tqdm_desc: Optional[str] = None,
**kwargs,
):
if tqdm_position is None:
tqdm_position = 0
if len(images) > batchsize:
return [
batch_run(
Expand All @@ -102,9 +106,17 @@ def batch_run(
post_args=post_args,
check=check,
wrap_errs=wrap_errs,
tqdm_position=tqdm_position + 1,
tqdm_desc=tqdm_desc,
stdout_write=stdout_write or partial(tqdm.write, file=sys.stdout),
stderr_write=stderr_write or partial(tqdm.write, file=sys.stderr),
**kwargs,
)
for i in range(0, len(images), batchsize)
for i in tqdm(
range(0, len(images), batchsize),
desc=f"Batch{' for {tqdm_desc}' if tqdm_desc else ''}",
position=tqdm_position,
)
]

if stdout_write is None:
Expand Down Expand Up @@ -166,6 +178,7 @@ def ect(
stdout_write: Optional[Callable] = None,
stderr_write: Optional[Callable] = None,
wrap_errs: Optional[list[Exception]] = None,
tqdm_position: Optional[int] = None,
):
if not images:
return
Expand All @@ -186,6 +199,7 @@ def ect(
stdout_write=stdout_write,
stderr_write=stderr_write,
wrap_errs=wrap_errs,
tqdm_position=tqdm_position,
)


Expand All @@ -209,6 +223,7 @@ def optipng(
stdout_write: Optional[Callable] = None,
stderr_write: Optional[Callable] = None,
wrap_errs: Optional[list[Exception]] = None,
tqdm_position: Optional[int] = None,
):
if not images:
return
Expand All @@ -225,6 +240,7 @@ def optipng(
stdout_write=stdout_write,
stderr_write=stderr_write,
wrap_errs=wrap_errs,
tqdm_position=tqdm_position,
)


Expand All @@ -241,6 +257,7 @@ def pngcrush(
stdout_write: Optional[Callable] = None,
stderr_write: Optional[Callable] = None,
wrap_errs: Optional[list[Exception]] = None,
tqdm_position: Optional[int] = None,
):
if not images:
return
Expand Down Expand Up @@ -276,6 +293,7 @@ def pngcrush(
stdout_write=stdout_write,
stderr_write=stderr_write,
wrap_errs=wrap_errs,
tqdm_position=tqdm_position,
)

# Replace original images with crushed images if they are smaller
Expand All @@ -296,7 +314,7 @@ def optimize(
tmpdir: Optional[Union[str, Path]] = None,
cleanup: Optional[bool] = None,
trim_printout: bool = False,
tqdm_position: Optional[int] = None,
tqdm_position: int = 0,
tqdm_leave: Optional[bool] = None,
stdout_write: Optional[Callable] = None,
stderr_write: Optional[Callable] = None,
Expand All @@ -316,6 +334,7 @@ def optimize(
stdout_write=partial(tqdm.write, file=sys.stdout),
stderr_write=partial(tqdm.write, file=sys.stderr),
wrap_errs=wrap_errs,
tqdm_position=tqdm_position + 1,
)
optipng(
*cur_images,
Expand All @@ -324,6 +343,7 @@ def optimize(
stdout_write=stdout_write,
stderr_write=stderr_write,
wrap_errs=wrap_errs,
tqdm_position=tqdm_position,
)
pngcrush(
*cur_images,
Expand All @@ -334,6 +354,7 @@ def optimize(
stdout_write=stdout_write,
stderr_write=stderr_write,
wrap_errs=wrap_errs,
tqdm_position=tqdm_position,
)
new_sizes = [Path(image).stat().st_size for image in cur_images]
cur_images = [
Expand Down

0 comments on commit e79f40a

Please sign in to comment.