JPEG thumbnails

Run thumbs.py to generate:
  - thumbnails
  - previews (enlarged and blurred thumbnails)
  - stripped image data
  - restored images out of stripped data

jpeg_headers

thumbs.py

#!/usr/bin/python3

from PIL import Image, ImageFilter, ExifTags
from io import BytesIO
from argparse import ArgumentParser
import sys, os, glob

def run_script(argv):
  parser = ArgumentParser()
  parser.add_argument('dir', help = 'Images dir', default = '')
  parser.add_argument('-t', '--thumb', help = 'Store thumbnails in "thumbs" dir', action = 'store_true')
  parser.add_argument('-p', '--preview', help = 'Store previews in "previews" dir', action = 'store_true')
  parser.add_argument('-d', '--data', help = 'Store stripped data in "data" dir', action = 'store_true')
  parser.add_argument('-r', '--restore', help = 'Restore images into "restored" dir', action = 'store_true')
  args = parser.parse_args()

  files = []
  for ext in ['*.jpg', '*.jpeg', '*.png', '*.gif']:
    files += glob.glob(os.path.join(args.dir, ext))

  for in_file in files:
    try:
      name = os.path.splitext(os.path.basename(in_file))[0] + '.jpg'
      print(name + '... ', end='')

      thumb = generate_thumbnail(in_file)

      data = strip_headers(thumb)
      print(str(len(data)) + ' bytes')

      if args.thumb: save_to_file(thumb, 'thumb/' + name)

      if args.data: save_to_file(BytesIO(data), 'data', name + '.dat')

      if args.restore: save_to_file(restore_image(data), 'restored', name)

      if args.preview: save_to_file(generate_preview(thumb), 'previews', name)

    except IOError as err:
      print('Error: {0}'.format(err))


# Generates thumbnail with common header
def generate_thumbnail(file):
  # Size and quality is picked in the way that provides the best quality in around 200-300 bytes.
  # Note: other values will have different common header.
  size = 36
  quality = 65

  with Image.open(file) as im:
    im = rotate_by_exif(im)
    im.thumbnail((size, size), resample = Image.BILINEAR) # Seems like BILINEAR has better quality
    im = remove_transparency(im)
    # Required to keep JPEG headers consistent (e.g. for greyscale images, images with transparency)
    im = im.convert(mode = 'RGB')

    out = BytesIO()
    # We must turn off optimization and use same qtables for all images ('web_low' has lowest quality)
    im.save(out, 'JPEG', quality = quality, optimize = False, qtables = 'web_low')
    return out

def remove_transparency(im, bg_colour=(255, 255, 255)):
  if im.mode in ('RGBA', 'LA') or (im.mode == 'P' and 'transparency' in im.info):
    background = Image.new('RGBA', im.size, bg_colour)
    return Image.alpha_composite(background, im.convert('RGBA'))
  else:
    return im

def rotate_by_exif(im):
  for orientation in ExifTags.TAGS.keys():
    if ExifTags.TAGS[orientation] == 'Orientation': break

  exif = im.getexif()
  if exif:
    if   exif[orientation] == 3: return im.rotate(180, expand = True)
    elif exif[orientation] == 6: return im.rotate(270, expand = True)
    elif exif[orientation] == 8: return im.rotate(90,  expand = True)
  return im


# Generates enlarged & blurred preview
def generate_preview(input):
  with Image.open(input) as im:
    aspect = im.width / im.height
    side = 256
    size = (side, int(side / aspect)) if aspect > 1 else (int(side * aspect), side)
    im = im.resize(size)
    #im = im.convert(mode = 'L') # Greyscale
    im = im.filter(ImageFilter.GaussianBlur(5))

    out = BytesIO()
    im.save(out, 'JPEG')
    return out


# Strips common headers from thumbnail and appends required metadata
def strip_headers(thumb):
  # Read common headers
  with open('jpeg_headers', 'rb') as file:
    headers = bytearray(file.read())

  thumb.seek(0)
  img = thumb.read()

  headers[164] = img[164] # Width
  headers[166] = img[166] # Height

  if headers != img[0:len(headers)]: raise IOError('Invalid common header')

  # Output is: [version, width, height, img_data]
  return bytes([1, img[164], img[166]]) + img[len(headers):len(img) - 2]

# Restores thumbnail from stipped data
def restore_image(data):
  # Read common headers
  with open('jpeg_headers', 'rb') as file:
    headers = bytearray(file.read())

  if data[0] != 1: raise IOError('Invalid data version')

  headers[164] = data[1] # Width
  headers[166] = data[2] # Height

  return BytesIO(headers + data[3:len(data)] + bytes.fromhex('ffd9'))


def save_to_file(bytes, dir, name):
  if not os.path.exists(dir): os.mkdir(dir)
  with open(os.path.join(dir, name), 'wb') as file:
    file.write(bytes.getbuffer())


run_script(sys.argv)