PageRenderTime 30ms CodeModel.GetById 1ms RepoModel.GetById 0ms app.codeStats 0ms

/go.py

https://github.com/davidwisch/va-pucks
Python | 199 lines | 145 code | 37 blank | 17 comment | 20 complexity | 6361dfe3b2c92c90d6f53adb940e53d9 MD5 | raw file
  1. import os
  2. import sys
  3. import time
  4. import datetime
  5. from lib import helpers
  6. from lib import analyze
  7. from lib import output
  8. from lib import cli
  9. BINARIES = "bin"
  10. MPLAYER = os.path.join(BINARIES, "mencoder", "mplayer.exe")
  11. PIL_INSTALLER = os.path.join(BINARIES, "installers", "PIL-1.1.7.win32-py2.6.exe")
  12. NUMPY_INSTALLER = os.path.join(BINARIES, "installers","numpy-1.3.0-win32-superpack-python2.6.exe")
  13. GNUPLOT = os.path.join(BINARIES, "gnuplot", "bin", "wgnuplot.exe")
  14. USER_OUTPUT_DIR = os.path.join(os.getenv('USERPROFILE'), 'Desktop')
  15. OUTPUT_DIR = "output"
  16. def check_deps():
  17. print "Checking Dependencies..."
  18. #check that mplayer.exe is around
  19. sys.stdout.write("Checking for mplayer...")
  20. if os.path.exists(MPLAYER):
  21. sys.stdout.write("FOUND\n")
  22. else:
  23. sys.stdout.write("MISSING\n")
  24. exit("ERROR: mplayer.exe was not found")
  25. sys.stdout.write("Checking for PIL...")
  26. try:
  27. import Image
  28. sys.stdout.write("FOUND\n")
  29. except:
  30. sys.stdout.write("MISSING\n")
  31. print 'Launching PIL installer in 3 seconds...'
  32. time.sleep(3)
  33. os.system(PIL_INSTALLER)
  34. exit("Exiting...relaunch when ready")
  35. sys.stdout.write("Checking for numpy...")
  36. try:
  37. import numpy
  38. sys.stdout.write("FOUND\n")
  39. except:
  40. sys.stdout.write("MISSING\n")
  41. print "Launching NumPy installer in 3 seconds..."
  42. time.sleep(3)
  43. os.system(NUMPY_INSTALLER)
  44. exit('Exiting...relaunch when ready')
  45. return True
  46. def init_dirs():
  47. if not os.path.exists(OUTPUT_DIR):
  48. os.mkdir(OUTPUT_DIR)
  49. if not os.path.exists(USER_OUTPUT_DIR):
  50. os.mkdir(USER_OUTPUT_DIR)
  51. if __name__ == "__main__":
  52. #Check dependencies
  53. check_deps()
  54. #parse the command line arguments
  55. options = cli.parse()
  56. #set number of bins
  57. if options.num_bins:
  58. NUM_BINS = options.num_bins
  59. else:
  60. NUM_BINS = 15
  61. #finish constructing USER_OUTPUT_DIR var
  62. folder_name = "%s.%s.%s.%s.%s.%s" % (
  63. str(datetime.date.today()),
  64. options.frame_interval,
  65. options.driver_power,
  66. options.table_angle,
  67. options.puck_mass,
  68. NUM_BINS
  69. )
  70. USER_OUTPUT_DIR = os.path.join(USER_OUTPUT_DIR, folder_name)
  71. init_dirs()
  72. output.output_dir_info(options, USER_OUTPUT_DIR, NUM_BINS)
  73. if not os.path.exists(options.video):
  74. exit("ERROR: Cannot find video file")
  75. #dump video to images
  76. CMD = "%s %s -nosound -speed 10.0 -vf framestep=%s -vo jpeg:outdir=%s:quality=80" % (MPLAYER, options.video, options.frame_interval, OUTPUT_DIR)
  77. if len(os.listdir(OUTPUT_DIR)) == 0:
  78. print "Dumping video to images... ( BE PATIENT )"
  79. time.sleep(2)
  80. output.dump_video(CMD)
  81. print "Dump complete"
  82. else:
  83. print "Output Directory not empty"
  84. delete = raw_input("Empty? [n/Y] ")
  85. if delete == "Y":
  86. helpers.cleanup(OUTPUT_DIR)
  87. output.dump_video(CMD)
  88. files = os.listdir(OUTPUT_DIR)
  89. print "Num Files:", len(files)
  90. bins = {}
  91. timesteps = []
  92. framesize = None
  93. failed_reads = 0
  94. total_pucks = 0
  95. for i in range(len(files)):
  96. filename = os.path.join(OUTPUT_DIR, files[i])
  97. puck_points, corner_points, frame_size = analyze.analyze_frame(filename, options)
  98. if not framesize:
  99. framesize = frame_size
  100. if len(corner_points) < 4:
  101. failed_reads += 1
  102. print "Unable to process frame %s/%s" % (i, len(files))
  103. continue
  104. upper_right, lower_right, lower_left, upper_left = analyze.find_corners(corner_points)
  105. if (upper_right is None or
  106. lower_right is None or
  107. lower_left is None or
  108. upper_left is None):
  109. failed_reads += 1
  110. print "Unable to process frame %s/%s" % (i, len(files))
  111. continue
  112. total_pucks += len(puck_points)
  113. #line between upper points
  114. upper_bins = helpers.get_upper_bins(upper_left, upper_right, NUM_BINS)
  115. #line between lower points
  116. lower_bins = helpers.get_lower_bins(lower_left, lower_right, NUM_BINS)
  117. #bin our data
  118. bins = analyze.bin_data(bins, puck_points, upper_bins, lower_bins)
  119. #save timestep
  120. timesteps.append(puck_points)
  121. if options.explicit:
  122. #write out bin points
  123. filename_bins = os.path.join(USER_OUTPUT_DIR, files[i]+'.bp.txt')
  124. output.output_bin_points(upper_bins, lower_bins, filename_bins)
  125. #writeout bin lines
  126. filename_lines = os.path.join(USER_OUTPUT_DIR, files[i]+'.bl.txt')
  127. output.output_bin_lines(upper_bins, lower_bins, filename_lines)
  128. #writeout puck points
  129. filename_points = os.path.join(USER_OUTPUT_DIR, files[i]+".pt.txt")
  130. output.output_2d_array(puck_points, filename_points)
  131. #writeout corner points
  132. filename_corners = os.path.join(USER_OUTPUT_DIR, files[i]+".cp.txt")
  133. output.output_2d_array(corner_points, filename_corners)
  134. #Create gnuplot file for plotting everything
  135. filename_gnuplot = os.path.join(USER_OUTPUT_DIR, files[i]+'.gnuplot')
  136. output.output_gnuplot(
  137. filename_gnuplot,
  138. filename_bins,
  139. filename_lines,
  140. filename_points,
  141. filename_corners
  142. )
  143. filename = os.path.join(USER_OUTPUT_DIR, "bincounts.txt")
  144. output.output_bin_data(bins, filename)
  145. #output heatmap
  146. heatmap = analyze.generate_heatmap(timesteps, framesize)
  147. output.output_heatmap(GNUPLOT, heatmap, USER_OUTPUT_DIR)
  148. #Output distribution plot
  149. distribution = analyze.generate_distribution(bins)
  150. dist_filename = output.output_distribution(GNUPLOT, distribution, USER_OUTPUT_DIR)
  151. print "Processed %s / %s" % (len(files) - failed_reads, len(files))
  152. try:
  153. print "Avg pucks detected/image: %s" % (total_pucks / (len(files)-failed_reads))
  154. except(ZeroDivisionError):
  155. pass
  156. print "Avg Puck Mass: %s" % (options.puck_mass)
  157. print "Table Angle: %s" % (options.table_angle)
  158. print "Driver Motor Power: %s" % (options.driver_power)
  159. print "Frame Interval: %s" % (options.frame_interval)
  160. print "Number of Bins: %s" % (NUM_BINS)
  161. print "Output Directory: %s" % (USER_OUTPUT_DIR)
  162. print "Raw Output is: %s" % (os.path.join(os.getcwd(), OUTPUT_DIR))