Source code for generate_frames

"""
 This code generates frames from CSV values that can be stiched together using FFMPEG 
 to animate pedestrian data.  This version produces an animation at 4x speed.
"""

print "Importing..."

#  Please ensure the following dependencies are installed before use:
import pylab
import numpy as np
import itertools
import sys, getopt
import operator
import collections


drawing_by_frame = []


#  
def generate_frames(argv):
	
	
	#  Some default values if nothing is provided in command line arguments.
	traces = 'bubble_pop_traces.csv'
	background = 'trails_480.png'
	
	#  Get command line arguments.
	#  -f specify a file name.  This code expects csv files in the format PedestrianID, X, Y, FrameNum
	#  -b specify a backgroun image.  Any format available to pylab is acceptable.	
	try:
		opts,args = getopt.getopt(argv, "f:b:")
	except getopt.GetoptError:
		print "Getopt Error"
		exit(2)
	
	for opt, arg in opts:
		if opt == "-f":
			traces = arg
		elif opt == "-b":
			background = arg

	#  Name each frame based on the filename		
	figure_name = traces.split("/")[-1].split(".")[-2]

	#  Load up csv file
	trace = np.loadtxt(traces, comments=';', delimiter=',')
	traces = itertools.groupby(trace, lambda x:x[0])
	
	#  These values should match those in pedestrian_tracking.py
	w,h=640,360
	border=20

	#  Some values from trail validation
	valid = 0
	avg_length = 0
	num_traces = 0

	#  Load up background image.
	background = pylab.imread(background)


	pylab.imshow(background)
	

	for id,t in traces:
		pts = np.array(list(t))
		invalid = False
	
		#  Validate Trails

		if (pts[0,1]>border and pts[0,1]<w-border) and (pts[0,2]>border and pts[0,2]<h-border):
			invalid = True

		if (pts[-1,1]>border and pts[-1,1]<w-border) and (pts[-1,2]>border and pts[-1,2]<h-border):
			invalid = True
		
		if len(pts) < 200:
			invalid = True
			
		if ((pts[0,2] > h-border) and (pts[0,1] > w/2-75 and pts[0,1] < w/2+75) or (pts[-1,2] > h-border) and (pts[-1,1] > w/2-75 and pts[-1,1] < w/2+75)):
			invalid = True

		
		#  For all valid trails, prepare them for generating animated trails by frame number	
		if not invalid:
					
			num_traces += 1
			avg_length += len(pts)
			
			#  Drawing colour for traces given as RGB
			colour = (0,0,1)
	
	
			for pt in pts:
				this_frame = [pt[3], pt[1], pt[2], pt[0]]				
				drawing_by_frame.append(this_frame)
					
			valid += 1
			
			x = np.clip(pts[:,1],0,w)
			y = np.clip(pts[:,2],0,h)
		
			
		
	print "Valid Trails: " , valid, " Average Length:" , avg_length/num_traces
	
	
	drawing_by_frame.sort()	
	
	last_frame = drawing_by_frame[-1][0]
	
	current_frame = drawing_by_frame[0][0]
	
	drawing_dict = collections.defaultdict(list)

	count = 0
	
	while len(drawing_by_frame) > 0:
	
		#print "Next Frame, " , current_frame
		pylab.imshow(background)
		
		while drawing_by_frame[0][0] == current_frame:
			list_one = drawing_by_frame.pop(0)
			
			x = drawing_dict[list_one[3]]
			x.append([list_one[1], list_one[2]])
			drawing_dict[list_one[3]] = x
			
		#  Adjust mod value here to adjust frame drawing frequency
		#  Draw stuff here
		if (current_frame % 10 ==0):
			print "Percentage Complete: " , (current_frame/last_frame)*100
			draw_dict(drawing_dict, w, h, border, figure_name, current_frame, count)

			count += 1
		
		pylab.clf()
		current_frame = drawing_by_frame[0][0]


def draw_dict(dict, w, h, border, figure_name, frame, count):
	for trace in dict:
		
		print trace
		pts = dict[trace]
	
		pylab.plot([p[0] for p in pts], [p[1] for p in pts],'-',color=(0,0,1),alpha=0.5, linewidth=2)	
		pylab.xlim(0,w)
		pylab.ylim(h,0)
		pylab.axis('off')
		pylab.subplots_adjust(0,0,1,1,0,0)
	pylab.savefig("Frames/" + figure_name + "_" + str(count).zfill(6) + '.png', dpi=150,bbox_inches='tight', pad_inches=0)	
	#pylab.savefig("Frames/" + 'frame' + str(int(frame)) + '.png', dpi=150,bbox_inches='tight', pad_inches=0)	
				

if __name__ == "__main__":

	print "Starting Frame Generation"
	generate_frames(sys.argv[1:])