# Installs the latest version of TFC compatible with the installed TF version.
!pip install tensorflow-compression~=$(pip show tensorflow | perl -p -0777 -e 's/.*Version: (\d\.\d).*/\1.0/sg')

# Downloads the 'models' directory from Github.
![[ -e /tfc ]] || git clone https://github.com/tensorflow/compression /tfc
%cd /tfc/models

# Checks if tfci.py is available.
import tfci


"""
Explanation: High-Fidelity Generative Image Compression
This colab can be used to compress images using HiFiC. This can also be achieved
by running tfci.py, as explained in the README.
Please visit hific.github.io for more information.
Setup Colab
End of explanation
"""


import tensorflow as tf

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

"""
Explanation: 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
End of explanation
"""


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')

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)



"""
Explanation: Imports and Definitions
End of explanation
"""


#@title Setup { vertical-output: false, run: "auto", display-mode: "form" }
#@markdown #### Custom Images
#@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", label:"HI"}

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)

#@markdown #### Select a model
#@markdown Different models target different bitrates.

model = 'hific-lo' #@param ["hific-lo", "hific-mi", "hific-hi"]



if 'upload_custom_images' not in locals():
  print('ERROR: Please run the previous cell!')
  # Setting defaults anyway.
  upload_custom_images = False
  model = 'hific-lo'

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

all_files = os.listdir(FILES_DIR)
print(f'Got the 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)

"""
Explanation: Load files
End of explanation
"""


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)

  compressed_path = os.path.join(OUT_DIR, f'{file_name}_{model}.tfci')
  output_path = os.path.join(OUT_DIR, f'{file_name}_{model}.png')
  
  if os.path.isfile(output_path):
    print('Exists already:', output_path)
    num_bytes = os.path.getsize(compressed_path)
    all_outputs.append(
      File(output_path, num_bytes,
           get_bpp(Image.open(full_path).size, num_bytes)))
    continue

  print('Compressing', file_name, 'with', model, '...')
  tfci.compress(model, full_path, compressed_path)
  num_bytes = os.path.getsize(compressed_path)
  print(f'Compressed to {num_bytes} bytes.')

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

print('All done!')

"""
Explanation: Compress images
End of explanation
"""


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/>')

"""
Explanation: Show output
End of explanation
"""


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) 

"""
Explanation: 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.
End of explanation
"""


