Skip to content

Commit

Permalink
some more options
Browse files Browse the repository at this point in the history
  • Loading branch information
lucidrains committed Jun 2, 2023
1 parent 1c2f18c commit 7ac6b3a
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 7 deletions.
33 changes: 27 additions & 6 deletions TPDNE_utils/tpdne.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,24 @@ def sample_image_and_save_repeatedly(
output_path: str = './out/random', # path to the output image, without extension (will be saved as webp)
*,
call_every_ms: int = 250, # how often to sample
tmp_dir = '/tmp', # to store temporary images, before symbolically linking to the output path
num_rotated_tmp_images = 10,
verbose: bool = True
tmp_dir: str = '/tmp', # to store temporary images, before symbolically linking to the output path
num_rotated_tmp_images: int = 10,
image_format: str = 'webp',
verbose: bool = True,
quality = 99,
generate_favicon: bool = True,
favicon_size: int = 32
):
assert 0 < quality <= 100
assert favicon_size in {16, 32}
assert image_format in {'webp', 'png', 'jpeg'}

tmp_dir = Path(tmp_dir)
output_path = Path(output_path).with_suffix('.webp')
output_path = Path(output_path)
assert output_path.suffix == '', 'output path suffix will be automatically determined by `image_format` keyword arg'

output_path = output_path.with_suffix(f'.{image_format}')

call_every_seconds = call_every_ms / 1000

assert tmp_dir.is_dir()
Expand All @@ -32,11 +44,20 @@ def sample_image_and_save_repeatedly(
image_tensor = fn()

tmp_image_index = (tmp_image_index + 1) % num_rotated_tmp_images
tmp_path = str(tmp_dir / f'{tmp_image_index}.webp')
tmp_path = str(tmp_dir / f'{tmp_image_index}.{image_format}')

Image.fromarray(image_tensor, 'RGB').save(tmp_path, format = 'webp')
pil_image = Image.fromarray(image_tensor, 'RGB')
pil_image.save(tmp_path, format = image_format, quality = quality)
os.system(f'ln -nfs {tmp_path} {output_path}')

if generate_favicon:
tmp_favicon_path = str(tmp_dir / f'favicon_{tmp_image_index}.png')
output_favicon_path = output_path.parents[0] / 'favicon.png'

small_pil_image = pil_image.resize((favicon_size, favicon_size))
small_pil_image.save(tmp_favicon_path)
os.system(f'ln -nfs {tmp_favicon_path} {output_favicon_path}')

elapsed = time() - start

if verbose:
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
setup(
name = 'TPDNE-utils',
packages = find_packages(exclude=[]),
version = '0.0.2',
version = '0.0.3',
license='MIT',
description = 'TPDNE',
author = 'Phil Wang',
Expand Down

0 comments on commit 7ac6b3a

Please sign in to comment.