jenc.py

from scipy.fftpack import dctn,idctn
from numpy import any, clip, floor, array, asarray, zeros, round, resize
from functools import partial

# Orthonormal DCT is required for JPEG quantisation else
# the default one is quite crap and useless (cannot
# reproduce values after transformation).
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

import logging

root = logging.RootLogger(logging.DEBUG)
root.addHandler(logging.StreamHandler())

def jpeglike_quant(imagein, Q=100):
    def qtab(Q, luma=True):
        """Generate DC quantization table from a Q-factor
        using the formula defined by the IJG.

        - Q: controls the coarseness of quantisation, the
        range is limited in (+24..+255].

        - luma: whether to generate quantisation table for
        luma or chroma channel. 
        """
        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]]) \
                if luma else array([
                [17, 18, 24, 47, 99, 99, 99, 99],
                [18, 21, 26, 66, 99, 99, 99, 99],
                [24, 26, 56, 99, 99, 99, 99, 99],
                [47, 66, 99, 99, 99, 99, 99, 99],
                [99, 99, 99, 99, 99, 99, 99, 99],
                [99, 99, 99, 99, 99, 99, 99, 99],
                [99, 99, 99, 99, 99, 99, 99, 99],
                [99, 99, 99, 99, 99, 99, 99, 99]])
        
        Q = clip(Q, 1, 100)
        S = 5000/Q if Q < 50 else 200 - 2*Q

        qstable = floor((base * S + 50) / 100)
        qstable[qstable == 0] = 1

        return qstable.astype('B')
    
    def qt_then_dqt(blk, qtab):
        assert blk.shape == (8,8), "not a valid 8x8 block."

        blk = round(dctn(blk.copy()) / qtab)
        return idctn(blk * qtab)
    
    sw, sh  = imagein.size
    imagein = resize(asarray(
                imagein.convert('YCbCr')).copy(),
              (ceilmod(sh, 8), ceilmod(sw, 8), 3))

    # Generate different quantisation tables
    # for luma and chroma quantisation.
    lqtab, cqtab = qtab(Q), qtab(Q, False)

    root.debug(
        "Luma QT: \n%s\n"
        "Chroma QT: \n%s" %
        (lqtab, cqtab))

    # 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[:2]
    for y in range(0, h, 8):
        for x in range(0, w, 8):
            imagein[y:y+8,x:x+8,0] = qt_then_dqt(imagein[y:y+8,x:x+8,0], lqtab)
            imagein[y:y+8,x:x+8,1] = qt_then_dqt(imagein[y:y+8,x:x+8,1], cqtab)
            imagein[y:y+8,x:x+8,2] = qt_then_dqt(imagein[y:y+8,x:x+8,2], cqtab)

    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 DC \n"
        "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)