|
|
|
import numpy as np
|
|
|
|
|
|
def quantize(arr, min_val, max_val, levels, dtype=np.int64):
|
|
"""Quantize an array of (-inf, inf) to [0, levels-1].
|
|
|
|
Args:
|
|
arr (ndarray): Input array.
|
|
min_val (scalar): Minimum value to be clipped.
|
|
max_val (scalar): Maximum value to be clipped.
|
|
levels (int): Quantization levels.
|
|
dtype (np.type): The type of the quantized array.
|
|
|
|
Returns:
|
|
tuple: Quantized array.
|
|
"""
|
|
if not (isinstance(levels, int) and levels > 1):
|
|
raise ValueError(
|
|
f'levels must be a positive integer, but got {levels}')
|
|
if min_val >= max_val:
|
|
raise ValueError(
|
|
f'min_val ({min_val}) must be smaller than max_val ({max_val})')
|
|
|
|
arr = np.clip(arr, min_val, max_val) - min_val
|
|
quantized_arr = np.minimum(
|
|
np.floor(levels * arr / (max_val - min_val)).astype(dtype), levels - 1)
|
|
|
|
return quantized_arr
|
|
|
|
|
|
def dequantize(arr, min_val, max_val, levels, dtype=np.float64):
|
|
"""Dequantize an array.
|
|
|
|
Args:
|
|
arr (ndarray): Input array.
|
|
min_val (scalar): Minimum value to be clipped.
|
|
max_val (scalar): Maximum value to be clipped.
|
|
levels (int): Quantization levels.
|
|
dtype (np.type): The type of the dequantized array.
|
|
|
|
Returns:
|
|
tuple: Dequantized array.
|
|
"""
|
|
if not (isinstance(levels, int) and levels > 1):
|
|
raise ValueError(
|
|
f'levels must be a positive integer, but got {levels}')
|
|
if min_val >= max_val:
|
|
raise ValueError(
|
|
f'min_val ({min_val}) must be smaller than max_val ({max_val})')
|
|
|
|
dequantized_arr = (arr + 0.5).astype(dtype) * (max_val -
|
|
min_val) / levels + min_val
|
|
|
|
return dequantized_arr
|
|
|
