Simple Nearest Neighbor Scaling Algorithm

nearest_scale.py

#
# Simple Nearest Neighbor Scaling Algorithm
#
# @author       Cameron Eure
# @license      Public Domain
#

import os, sys, copy

try:
    import png
    HAVE_PNG = True
except:
    HAVE_PNG = False
    print('PNG Library not Found... Quiting')
    sys.exit(1)

def chunks(source, n):
    """ Split list into chunks of n size """
    for i in xrange(0, len(source), n):
        yield source[i:i+n]

class Image(object):
    """ Class which represents an image """
    
    def __init__(self, width, height, data=[], pitch=None):
        if pitch is None:
            pitch = width
        
        self.width = width
        self.height = height
        self.pitch = pitch
        self.data = []
        
        if len(data) == 0:
            self.data = [[0, 0, 0] for i in range(pitch * height)]
        else:
            self.data = data
    
    def get_pixel_idx(self, x, y):
        return (y * self.pitch) + x
    
    def get_pixel(self, x, y):
        return self.data[self.get_pixel_idx(x, y)]
    
    def set_pixel(self, x, y, rgb):
        self.data[self.get_pixel_idx(x, y)] = [i for i in rgb]
    
    def get_box(self, center_x, center_y, size=3):
        """
        Returns a matrix of pixels of (size)x(size), with pixels from starting point
        (center_x, center_y). If pixel is out of range, None is used in place of a
        list of RGB values.
        """
    
        median = size / 2
        start = -(median)
        end = median+1
        box = [[None for a in range(size)] for b in range(size)]
        
        max_width = self.width - 1
        max_height = self.height - 1
        
        ps = 0
        for yo in range(start, end):
            px = 0
            for xo in range(start, end):
                xt = center_x + xo
                yt = center_y + yo
                
                if (not (
                        xt < 0 or
                        xt > max_width or
                        yt < 0 or
                        yt > max_height)):
                    box[ps][px] = self.data[self.get_pixel_idx(xt, yt)]
                
                px += 1
            ps += 1
        
        return box
    
    def savePNG(self, filename):
        if not HAVE_PNG:
            raise StandardError('SavePNG() not available without PyPNG')
    
        tmp = []
        for i in self.data:
            tmp.extend(i)
        
        tmp = list(chunks(tmp, self.width * 3))
        out = png.Writer(width=self.width, height=self.height)
        out.write(open(filename, 'w'), tmp)
        del out
        del tmp
    
    @staticmethod
    def fromPNG(filename):
        if not HAVE_PNG:
            raise StandardError('SavePNG() not available without PyPNG')
    
        r = png.Reader(filename=filename)
        r.read()
        width, height, pixels, meta = r.asRGB8()
        data = []
        
        for row in pixels:
            data.extend(list(chunks(row, 3)))
        
        del r
        del pixels
        return Image(
            width=width,
            height=height,
            data=data)

def main():
    if len(sys.argv) < 5:
        print('Usage: %s <input.png> <output.png> <scale width> <scale height>' % (sys.argv[0]))
        sys.exit(1)
    
    # Input / Output file paths
    file_in = sys.argv[1]
    file_out = sys.argv[2]
    
    # New Width
    try:
        scale_width = int(sys.argv[3])
    except ValueError:
        print('Width must be an integer')
        sys.exit(1)
    
    # New Height
    try:
        scale_height = int(sys.argv[4])
    except ValueError:
        print('Height must be an integer')
        sys.exit(1)
    
    # Load input file
    try:
        image_in = Image.fromPNG(file_in)
    except:
        print('Failed to load PNG image %s' % (file_in))
        sys.exit(1)
    
    # Create output buffer
    image_out = Image(width=scale_width, height=scale_height)
    
    #
    # Mathmatically, this is the same as using floating point. We just move the
    # input width/height to a different scale (the "<< 16" part), so we can use
    # integers instead.
    #
    y2_mult = (image_in.height << 16) / image_out.height + 1
    x2_mult = (image_in.width << 16) / image_out.width + 1
    
    #
    # Logic: For each output pixel, find the corresponding input pixel. By dividing
    # the input width/height by the output width/height, we can determine the x/y
    # coordinate of the pixel that most closely matches the output pixel.
    #
    for y in range(image_out.height):
        # Get input Y
        y2 = (y * y2_mult) >> 16
        
        for x in range(image_out.width):
            # Get input X
            x2 = (x * x2_mult) >> 16
            
            # Set output pixel
            image_out.set_pixel(x, y, image_in.get_pixel(x2, y2))

    try:
        image_out.savePNG(file_out)
    except:
        print('Failed to save output to %s' % (file_out))
        sys.exit(1)

if __name__ == '__main__':
    main()