Setup Colab



In [ ]:

    
%tensorflow_version 1.x
!pip install tensorflow-compression
![[ -e tfc ]] || git clone https://github.com/tensorflow/compression tfc
%cd tfc/models
import tfci  # Check if tfci.py is available.

Enabling GPU

GPU should be enabled for this colab. If the next cell prints a warning, do the following:

Navigate to Edit→Notebook Settings
select GPU from the Hardware Accelerator drop-down



In [ ]:

    
import tensorflow as tf

if not tf.test.is_gpu_available():
  print('WARNING: No GPU found. Might be slow!')
else:
  print('Found GPU.')

Imports and Definitions



In [ ]:

    
import os
import zipfile
from google.colab import files
import collections
from PIL import Image
from IPython.display import Image as DisplayImage
from IPython.display import Javascript
from IPython.core.display import display, HTML
import tfci
import urllib.request

tf.get_logger().setLevel('WARN')  # Only show Warnings

FILES_DIR = '/content/files'
OUT_DIR = '/content/out'
DEFAULT_IMAGE_URL = ('https://storage.googleapis.com/hific/clic2020/'
                     'images/originals/ad249bba099568403dc6b97bc37f8d74.png')

MODEL = 'hific-lo'
TMP_OUT = 'out.tfci'

os.makedirs(FILES_DIR, exist_ok=True)
os.makedirs(OUT_DIR, exist_ok=True)

File = collections.namedtuple('File', ['full_path', 'num_bytes', 'bpp'])

def print_html(html):
  display(HTML(html + '<br/>'))

def make_cell_large():
  display(Javascript(
      '''google.colab.output.setIframeHeight(0, true, {maxHeight: 5000})'''))

def get_default_image(output_dir):
  output_path = os.path.join(output_dir, os.path.basename(DEFAULT_IMAGE_URL))
  print('Downloading', DEFAULT_IMAGE_URL, '\n->', output_path)
  urllib.request.urlretrieve(DEFAULT_IMAGE_URL, output_path)


print('Caching model...')
tfci.import_metagraph(MODEL)
print('Done')

Load files



In [ ]:

    
#@title Loading Images { vertical-output: false, run: "auto", display-mode: "form" }
#@markdown Tick the following if you want to upload your own images to compress.
#@markdown Otherwise, a default image will be used.
#@markdown 
#@markdown **Note**: We support JPG and PNG (without alpha channels).
#@markdown 

upload_custom_images = False #@param {type:"boolean"}

if upload_custom_images:
  uploaded = files.upload()
  for name, content in uploaded.items():
    with open(os.path.join(FILES_DIR, name), 'wb') as fout:
      print('Writing', name, '...')
      fout.write(content)



In [ ]:

    
all_files = os.listdir(FILES_DIR)
if not upload_custom_images or not all_files:
  print('Downloading default...')
  get_default_image(FILES_DIR)
  print()

all_files = os.listdir(FILES_DIR)
print(f'Got following files ({len(all_files)}):')

for file_name in all_files:
  img = Image.open(os.path.join(FILES_DIR, file_name))
  w, h = img.size
  img = img.resize((w // 15, h // 15))
  print('- ' + file_name + ':')
  display(img)

Compress images



In [ ]:

    
SUPPORTED_EXT = {'.png', '.jpg'}

all_files = os.listdir(FILES_DIR)
if not all_files:
  raise ValueError("Please upload images!")

def get_bpp(image_dimensions, num_bytes):
  w, h = image_dimensions
  return num_bytes * 8 / (w * h)

def has_alpha(img_p):
  im = Image.open(img_p)
  return im.mode == 'RGBA'

all_outputs = []
for file_name in all_files:
  if os.path.isdir(file_name):
    continue
  if not any(file_name.endswith(ext) for ext in SUPPORTED_EXT):
    print('Skipping', file_name, '...')
    continue
  full_path = os.path.join(FILES_DIR, file_name)
  if has_alpha(full_path):
    print('Skipping because of alpha channel:', file_name)
    continue
  file_name, _ = os.path.splitext(file_name)
  output_path = os.path.join(OUT_DIR, f'{file_name}.png')
  if os.path.isfile(output_path):
    print('Skipping', output_path, '-- exists already.')
    continue

  print('Compressing', file_name, '...')
  tfci.compress(MODEL, full_path, TMP_OUT)
  num_bytes = os.path.getsize(TMP_OUT)

  print('Decompressing...')
  tfci.decompress(TMP_OUT, output_path)
  
  all_outputs.append(
      File(output_path, num_bytes,
           get_bpp(Image.open(full_path).size, num_bytes)))

print('All done!')

Show output



In [ ]:

    
make_cell_large()  # Larger output window.

for file in all_outputs:
  print_html('<hr/>')
  print(f'Showing {file.full_path} | {file.num_bytes//1000}kB | {file.bpp:.4f}bpp')
  display(Image.open(file.full_path))
  print_html('<hr/>')

Download all compressed images.

To download all images, run the following cell.

You can also use the Files tab on the left to manually select images.

Note: the images are saved as PNGs and thus very large. The bitrate used by HiFiC is given in the name.



In [ ]:

    
ZIP = '/content/images.zip'

with zipfile.ZipFile(ZIP, 'w') as zf:
  for f in all_outputs:
    path_with_bpp = f.full_path.replace('.png', f'-{f.bpp:.3f}bpp.png')
    zf.write(f.full_path, os.path.basename(path_with_bpp))

files.download(ZIP)