jenc.py

# DCT is required for quantising coefficients, this always
# performs a orthonormal DCT and IDCT.
from scipy.fftpack import dctn,idctn
from numpy import any, clip, floor, array, asarray, zeros, round, resize
from functools import partial

# Explicitly specify that Type-II DCT is required.
dctn  = partial(dctn,  type=2, norm='ortho')
idctn = partial(idctn, type=2, norm='ortho')

ceilmod = lambda n,mod: n if n % mod == 0 else n + mod - n % mod

# Generate a quantisation table based on the IJG formula
# provided, this is not considered great but still okay.
def genqtab(Q):
    """Generate 8x8 quantisation table from a given Q factor
    in (24..100] range (24=worst, 100=best)."""

    # base table for quantization of values.
    base = array([
            [16, 11, 10, 16,  24,  40,  51,  61],
            [12, 12, 14, 19,  26,  58,  60,  55],
            [14, 13, 16, 24,  40,  57,  69,  56],
            [14, 17, 22, 29,  51,  87,  80,  62],
            [18, 22, 37, 56,  68, 109, 103,  77],
            [24, 35, 55, 64,  81, 104, 113,  92],
            [49, 64, 78, 87, 103, 121, 120, 101],
            [72, 92, 95, 98, 112, 100, 103,  99]])
    
    # find the S factor from the Q factor.
    Q = clip(Q, 24, 100)
    S = 5000/Q if Q < 50 else 200 - 2*Q

    # build the quantisation table from the S factor by
    # applying some math on Annex-K table.
    qstable = floor((base * S + 50) / 100)
    qstable[qstable == 0] = 1

    return qstable.astype('B')

def jpeglike_quant(imagein, Q=100):
    sw, sh  = imagein.size
    imagein = resize(asarray(
                imagein.convert('YCbCr')).copy(),
              (ceilmod(sh, 8), ceilmod(sw, 8), 3))

    # generate quantization table for a given factor.
    qtab = genqtab(abs(int(Q)))

    # convert 8-bit pixel values into dct coefficients and quantize
    # them the generated quantization table and de-quantize them
    # immdieately and replace at the same position.
    h,w,_ = imagein.shape
    for y in range(0, h, 8):
        for x in range(0, w, 8):
            for c in range(3):
                qsimg  = round(dctn(imagein[y:y+8,x:x+8,c]) / qtab)
                imagein[y:y+8,x:x+8,c] = idctn(qsimg * qtab)

    return Image.fromarray(
        resize(imagein, (sh, sw, 3)), 
        mode='YCbCr')
    

if __name__ == "__main__":
    from PIL import Image
    from argparse import ArgumentParser, FileType, RawTextHelpFormatter

    parser = ArgumentParser(__file__,
        description=
        "This app is used to simulate JPEG encoding by doing\n"
        "actual DCT quantisation and dequantisation of coefficients.\n"
        "This is made for getting those artefacts generated from\n"
        "DCT coefficient quantisation.\n",
        formatter_class=RawTextHelpFormatter)

    parser.add_argument('-q',
        type=int,
        default=95,
        help="Q-factor to use for generating a quantisation table.\n"
             "(lesser values result in a coarse table).")

    parser.add_argument('FILE',
        type=FileType('rb'),
        help="Input file to process through PIL.")
    
    parser.add_argument('OUTPUT',
        type=FileType('wb'),
        help="Output file to dump the processed image.")

    argv = parser.parse_args()

    jpeglike_quant(Image.open(argv.FILE), argv.q) \
        .convert('RGB') \
        .save(argv.OUTPUT)