Contrast correct an hdf5 file.

contrast_correct.py

import numpy as np
import h5py
from tqdm import tqdm

from cloudvolume import view

def find_section_clamping_values(zlevel, lowerfract, upperfract):
  filtered = np.copy(zlevel)

  # remove pure black from frequency counts as
  # it has no information in our images
  filtered[0] = 0

  cdf = np.zeros(shape=(len(filtered),), dtype=np.uint64)
  cdf[0] = filtered[0]
  for i in range(1, len(filtered)):
    cdf[i] = cdf[i - 1] + filtered[i]

  total = cdf[-1]

  if total == 0:
    return (0, 0)

  lower = 0
  for i, val in enumerate(cdf):
    if float(val) / float(total) > lowerfract:
      break
    lower = i

  upper = 0
  for i, val in enumerate(cdf):
    if float(val) / float(total) > upperfract:
      break
    upper = i

  return (lower, upper)

CLIP_FRACTION = 0.01
FILENAME = 'val_image'

with h5py.File('./{}.h5'.format(FILENAME), 'r') as f:
  img = f['main'][:]

num_z = img.shape[2]
area = img.shape[0] * img.shape[1]
dtype = img.dtype

nbits = np.dtype(img.dtype).itemsize * 8
num_vals = 2 ** nbits
maxval = float(2 ** nbits - 1)

for z in tqdm(range(num_z)):
  # compute luminance levels
  levels = np.zeros(shape=num_vals, dtype=np.uint64)
  img2d = img[:,:,z]
  cts = np.bincount(img2d.reshape(( area )))
  levels[0:len(cts)] += cts.astype(np.uint64)

  # perform contrast correction
  (lower, upper) = find_section_clamping_values(levels, CLIP_FRACTION, 1 - CLIP_FRACTION)
  if lower == upper:
    continue

  img2d = img2d.astype(np.float32)
  img2d = (img2d - float(lower)) * (maxval / (float(upper) - float(lower)))
  img2d = np.round(img2d)
  img2d = np.clip(img2d, 1, maxval)
  img[:, :, z] = img2d.astype(dtype)

with h5py.File('./{}_contrast_corrected_{}.h5'.format(FILENAME, CLIP_FRACTION), 'w') as f:
  f.create_dataset('main', data=img)