Tracking/generate_frames.py at master · Grad-CSE2016/Tracking · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

"""
 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:])