Testing JPNG file format

jpng.py

import os, os.path
import numpy as np
import scipy
import scipy.misc
import scipy.ndimage.morphology as morphology
import subprocess

def printarr(arr):
    print "\n".join([''.join(map(str,row)) for row in arr])

#-----------------------------------------------------------------------------------------
# Make a chequered flag image
#
base = np.array([1]*8 + [0]*8)
stripe1 = np.tile(base,4*8).reshape((8,64))
stripe2 = (1-stripe1)
stripe3 = np.vstack((stripe1,stripe2))
img = np.tile(stripe3.T, 4)
scipy.misc.imsave("img.png", img)

print "Chequered flag".center(120, "_")
printarr(img[::4,::4])
print "img.shape", img.shape
print "img png size", os.path.getsize("img.png")

#-----------------------------------------------------------------------------------------
# Now move a column from the left to the right -- chaos!
#
quality = 85
for n in range(0,9):
    shifted_img = np.hstack((img[:,n:], img[:,:n]))
    scipy.misc.imsave("shifted_img.png", shifted_img)
    subprocess.call(["convert", "-quality", str(quality), "shifted_img.png", "img.jpeg"])
    sz = os.path.getsize("img.jpeg")
    print "shifted %d column: jpeg size" % n, sz, "#"*int(sz/20.0)

#-----------------------------------------------------------------------------------------
# Resulting image should be at least as close to the original image as 
# the original jpeg, plus the two images together should be smaller in size than the 
# baseline jpeg
#

def find_low_freq(img, N):
    # Set to True any square that is below a certain frequency threshold
    fg_squares = np.zeros(img.shape[:2], np.bool)
    for y in range(0, len(img), N):
        for x in range(0, len(img[0]), N):
            eight = img[y:y+N, x:x+N] / 255.
            freqx = np.abs(eight[:-1,:] - eight[1:,:])
            freqy = np.abs(eight[:,:-1] - eight[:,1:])
            #printarr((.5+eight[:,:,0]).astype(np.int))
            tfreq = freqx.sum() + freqy.sum()
            if tfreq < thresh: # low freq = fg = png
                fg_squares[y:y+N,x:x+N] = True
            else:
                fg_squares[y:y+N,x:x+N] = False
    return fg_squares

def find_inaccurate_blocks(img, jpeg, thresh, N):
    """Compare img to jpeg version and find squares that are different"""
    fg_squares = np.zeros(img.shape[:2], np.bool)
    for y in range(0, len(img), N):
        for x in range(0, len(img[0]), N):
            eight1 = img[y:y+N, x:x+N] / 255.
            eight2 = jpeg[y:y+N, x:x+N] / 255.
            diff = np.abs(eight1 - eight2).sum()
            if diff > thresh:
                fg_squares[y:y+N,x:x+N] = True
    return fg_squares

def make_jpng(original_fname, thresh, N=8, quality=20):
    root, _ext = os.path.splitext(original_fname)
    original = scipy.misc.imread(original_fname)
    subprocess.call(["convert", "-quality", str(quality), original_fname, root+".jpeg"])
    orig_jpeg = scipy.misc.imread(root+".jpeg")
        
    size_original = os.path.getsize(original_fname)
    size_jpeg = os.path.getsize(root+".jpeg")
    
    fg_squares = find_inaccurate_blocks(original, orig_jpeg, thresh, N)
    
    #-------------------------------------------------------------------------------------
    # mask_bg is the jpeg of the image
    mask_bg = np.copy(original)
    fg_squares_shrunk = morphology.binary_erosion(fg_squares, structure=np.ones((9,9)))
    mask_bg[fg_squares_shrunk] = np.array([255,255,255])
    
    #---------------------------
    # mask_fg is transparent png
    mask_fg = np.copy(original)
    alpha = np.ones(mask_fg.shape[:2]) * 255
    mask_fg = np.dstack((mask_fg, alpha))
    mask_fg[~fg_squares] = np.array([255,255,255,0])
    
    # Save masked bg and fg as png, then convert to jpeg using imagemagick for consistency
    scipy.misc.imsave("mask_bg_img_{thresh:.3g}.png".format(thresh=thresh), mask_bg)
    scipy.misc.imsave("mask_fg_img_{thresh:.3g}.png".format(thresh=thresh), mask_fg)    
    subprocess.call(["convert", "-quality", str(quality), 
                     "mask_bg_img_{thresh:.3g}.png".format(thresh=thresh), 
                     "mask_bg_img_{thresh:.3g}.jpeg".format(thresh=thresh)])
    new_jpeg = scipy.misc.imread("mask_bg_img_{thresh:.3g}.jpeg".format(thresh=thresh))
    
    print "png equal?", np.array_equal(original, mask_bg)    
    print "jpeg equal?", np.array_equal(orig_jpeg, new_jpeg)
    
    # Get all image sizes
    size_mask_bg = os.path.getsize("mask_bg_img_{thresh:.3g}.jpeg".format(thresh=thresh))
    size_mask_fg = os.path.getsize("mask_fg_img_{thresh:.3g}.png".format(thresh=thresh))
    size_masked = size_mask_bg + size_mask_fg
    size_masked_str = "%sk (jpeg:%sk, png:%sk)" % (int(size_masked/1000.), 
        int(size_mask_bg/1000.), int(size_mask_fg/1000.))
    
    print str(thresh).ljust(3), "orig", size_original, "jpeg", size_jpeg, "new", 
    print size_masked, "(bg", size_mask_bg, "fg", size_mask_fg, ")"
    
    #-------------------------------------------------------------------------------------
    # Calculate distances from original(usually png) to regular jpeg and jpng
    #
    # Distance 1, a simple jpeg
    distance_jpeg = np.abs(orig_jpeg.astype(np.int) - original.astype(np.int)).sum()
    
    # Distance 2, a jpng image
    mask_fg2 = np.copy(mask_fg)[:,:,:3] # skip alpha channel
    mask_bg2 = np.copy(new_jpeg)
    mask_bg2[fg_squares] = np.array([0,0,0])
    mask_fg2[~fg_squares] = np.array([0,0,0])
    
    distance_masked = np.abs(mask_bg2.astype(np.int) + mask_fg2.astype(np.int)
                             - original.astype(np.int)).sum()
    
    return {"size_original":size_original, "size_jpeg":size_jpeg, 
            "size_masked":size_masked, "size_masked_str":size_masked_str,
            "distance_jpeg":distance_jpeg, "distance_masked":distance_masked}
 
def make_html(thresh, maskinfo):
    html = """<td>
    <img style="background:url(mask_bg_img_{thresh:.3g}.jpeg)" src="mask_fg_img_{thresh:.3g}.png" />
    <table border=1 style="border-collapse:true;border:solid 1px #ccc">
    <tr><td>Threshold: {thresh}</td>                 <td>Size of original: {size_original}k</td></tr>
    <tr><td>Error jpeg: {distance_jpeg}k</td>  <td>Size of jpeg: {size_jpeg}k</td></tr>
    <tr><td>Error jpng: {distance_masked}k</td><td>Size of jpng: {size_masked}</td></tr>
    </table>
    
    <table>
    <tr><td>JPEG</td><td>PNG</td></tr>
    <tr><td><img style="width:100%" src="mask_bg_img_{thresh:.3g}.jpeg" /></td>
    <td><img style="width:100%" src="mask_fg_img_{thresh:.3g}.png" /></td></tr>
    </table>
    """.format(thresh=thresh, size_original=int(maskinfo['size_original']/1000.),
               size_jpeg=int(maskinfo['size_jpeg']/1000.), size_masked=maskinfo['size_masked_str'],
               distance_jpeg=int(maskinfo['distance_jpeg']/1000.), 
               distance_masked=int(maskinfo['distance_masked']/1000.))
    return html

if __name__=="__main__":
    original_fname = "clooney3.png"
    html = []
    for _thresh in range(0,61,10) + [np.inf]:
        thresh = _thresh/10.
        maskinfo = make_jpng(original_fname, thresh)
        html.append(make_html(thresh, maskinfo))
    
    with open("index.html",'w') as out:
        out.write("<html><body><table><tr>" + '\n'.join(html) + "</tr></table></body></html>")