pixeler.py

#!/usr/bin/python

#
#	Title: Pixeler.py
#
#	Author: Dave "pugilist"
#
#	Date: 2014-12-01
#
#	License: CC BY-SA 4.0
#		https://creativecommons.org/licenses/by-sa/4.0/
#
#	Version: 0.5
#
#	Description:
#		A few silly tools for image modification.  
#		This is somewhat of a developer tool and has some serious limitations.
#		If you don't like that, go download GIMP or MS Paint or something. :)
#
#

from PIL import Image
import argparse
import random

def modulation_method(string):
	#this method defines the acceptable types of modulation, user input is checked here and if no match is found, an exception is raised.
	methods = ["modup","moddown","random","crazy"]

	if not string in methods:
		raise argparse.ArgumentTypeError("Error, unrecognised input: %s.\nAcceptable methods are: %s"%(string,methods))
	else:
		return string

def open_input_file(input_fh):	
	try:
		return Image.open(input_fh)
	except:
		print "[!] Cannot open input file, bad perms or path?"

def open_output_file(output_fh):
	# not sure if we're going to use this at all...
	pass	

def write_things(col_pointer,row_pointer,chunk_size,avg_rgb,canvas):
	for x in range(chunk_size):
		for y in range(chunk_size):
			# try to write the pre-defined color to all pixels, if the pixels are out of range, continue with the loop.
			try:
				canvas[col_pointer+x,row_pointer+y] = avg_rgb	
			except:
				continue

def get_rgb_avg(file_to_parse,chunk_size):
	# create an image object for use later
	img_obj 		= file_to_parse.load()

	# grab the dimensions of our input image
	(width,height) 		= file_to_parse.size

	# set up some vars before we iterate through our image
	row_pointer 			= 0
	col_pointer 			= 0
	total_pixels_processed  = 0

	# start iterating through the picture's rows 
	while True:
		# reset the column pointer 
		col_pointer = 0

		# iterate across the columns for the given row
		while True:
		
			if args.verbosity > 1: print "	[+] working on (" + str(col_pointer) + ", " + str(row_pointer) + ")"

			counter = 0
			r_total = 0
			g_total = 0
			b_total = 0

			for x in range(chunk_size):

				if (col_pointer+x) < width:

					for y in range(chunk_size):
						if (row_pointer+y) < height:
							if args.verbosity > 2: print "		[+] x: " + str(col_pointer+x) + ", y: " + str(row_pointer+y)

							#add exception handling for black and white images that only have one value
							try:
								(r,g,b)	= img_obj[col_pointer+x,row_pointer+y]
							except TypeError:
								r=g=b   = img_obj[col_pointer+x,row_pointer+y]
							r_total 		+= r
							g_total 		+= g
							b_total 		+= b
							counter 		+= 1
							total_pixels_processed  += 1

			# create a tuple with our chunk's average rgb value
			yield { 'avg_rgb':( r_total/counter,g_total/counter,b_total/counter ), 'col_pointer': col_pointer, 'row_pointer': row_pointer }	

			# iterate to the next chunk of the image
			col_pointer += chunk_size
			
			# of the pointer is greater than the width of the picture, start with the next row.
			if col_pointer >= width:
				break

		# iterate to the next row
		row_pointer += chunk_size
		
		# if we're at the end of the image, break out of the loop
		if row_pointer >= height:
			break

def modulate_rgb(input_rgb_values,embed_rgb_values,modulation_method):

	( r_i, g_i, b_i ) = input_rgb_values
	( r_e, g_e, b_e ) = embed_rgb_values

	'''
	# prototyped code for modulations
	for rgb_value in rgb_values:
		#do stuff
		if some_value:
			do some modulation:
	
	return (r,g,b)
	'''
	if modulation_method == "modup":
			return ( abs( r_i + r_e ), abs( g_i + g_e ), abs( b_i + b_e ) )

	elif modulation_method == "moddown":
			return ( abs( r_i - r_e ), abs( g_i - g_e ), abs( b_i - b_e ) )

	elif modulation_method == "random": 
		if random.getrandbits(1):
			return ( abs( r_i + r_e ), abs( g_i + g_e ), abs( b_i + b_e ) )
		else:
			return ( abs( r_i - r_e ), abs( g_i - g_e ), abs( b_i - b_e ) )

	elif modulation_method == "crazy":

		if random.getrandbits(1):
			r_n = abs( r_i + r_e )
		else:
			r_n = abs( r_i - r_e )

		if random.getrandbits(1):
			g_n = abs( g_i + g_e )
		else:
			g_n = abs( g_i - g_e )

		if random.getrandbits(1):
			b_n = abs( b_i - b_e )
		else:
			b_n = abs( b_i + b_e )
		
		return ( r_n, g_n, b_n )



def embed_file(args,canvas):
	
	# create instances of our generator function
	input_file_object = get_rgb_avg(args.input_fh,args.chunk_size)
	embed_file_object = get_rgb_avg(args.embed_fh,args.chunk_size)
	
	# for each chunk in our input file
	for input_file_chunk in input_file_object:
		# also get each chunk in our embed file.
		embed_file_chunk = embed_file_object.next()

		# modulate the rgb vals depending on the method we specify
		( r_n, g_n, b_n ) = modulate_rgb(input_file_chunk['avg_rgb'], embed_file_chunk['avg_rgb'], args.mod_method)

		# write our averaged chunks to the new image
		write_things(input_file_chunk['col_pointer'],input_file_chunk['row_pointer'],args.chunk_size,(r_n,g_n,b_n),canvas)	



def pixelate(args,canvas):
	
	# create instance of our generator function
	generator_instance = get_rgb_avg(args.input_fh,args.chunk_size)

	for items in generator_instance:
		avg_rgb = items['avg_rgb']

		if args.verbosity > 1: print "	[+] Avg RGB values: " + str(avg_rgb[0]) + ", " + str(avg_rgb[1]) + ", " + str(avg_rgb[2])	

		# write our averaged chunks to the new image
		write_things(items['col_pointer'],items['row_pointer'],args.chunk_size,items['avg_rgb'],canvas)	

	if args.verbosity > 2: print "[+] Total pixels processed: " + str(total_pixels_processed)




def main(args):
	# create output image so we can write our output to it
	output_img 		= Image.new( 'RGB', args.input_fh.size )
	canvas 	= output_img.load() 

	random.seed()
	

	if args.embed_fh:
		if args.verbosity > 0: print "[+] Starting Embedding process"
		embed_file(args,canvas)
	else:
		if args.verbosity > 0: print "[+] Starting Pixelating process"
		pixelate(args,canvas)

	if args.verbosity > 0: print "[+] Writing output file"
	
	#write the output file
	output_img.save(args.output_fh,"PNG")



if __name__ == "__main__":

	parser = argparse.ArgumentParser(description='Process all the pictures with the stuffs')
	parser.add_argument("input_fh",metavar="infile",type=open_input_file,help="the file that you want to parse")
	parser.add_argument("-o", "--output-file",metavar="outfile",dest="output_fh",type=str,required=True,help="the name of the output file")
	parser.add_argument("-e", "--embed",metavar="file_to_embed",dest="embed_fh",type=open_input_file,help="a file to embed in the input_fh")
	parser.add_argument("-m",metavar="method",dest="mod_method",type=modulation_method,help="Specify how pixels are modulated")
	parser.add_argument("-c","--chunk-size",metavar="chunk_size",type=int,default=10,help="the chunk size in pixels (default 10)")
	parser.add_argument("-v","--verbosity",action="count",help="add verbosity to the output")

	args = parser.parse_args()

	if args.embed_fh and not args.mod_method:
		# I coded this this way because I'm lazy, sue me.
		print "You need to specify -m with -e"
		exit(1)

	main(args)