frame counting bug fixed

ml4a/utils/video.py

@@ -1,5 +1,6 @@
 from PIL import Image, ImageStat, ImageEnhance
 import sys, argparse
+import glob
 import os
 from os import listdir
 from os.path import isfile, isdir, join
@@ -15,12 +16,12 @@
 
 
 def generate_video2(frames_path, output_path):
- numframes = len([f for f in listdir(frames_path) if isfile(join(frames_path, f))])
+ numframes = len(glob.glob('%s/f*.png'%frames_path))
  log('loading %d images from %s' % (numframes, frames_path))
  images = [ Image.open('%s/f%05d.png'%(frames_path, t+1)) for t in range(numframes) ]
  brightness = [ ImageStat.Stat(img.convert('L')).mean[0] for img in images ]
  avg_brightness = np.mean(brightness)
- os.system('mkdir %s/temp' % frames_path)
+ os.system('mkdir "%s/temp"' % frames_path)
  for i, img in tqdm(enumerate(images)):
  mult = avg_brightness / brightness[i]
  source = img.split()
@@ -32,23 +33,23 @@ def generate_video2(frames_path, output_path):
  img2 = ImageEnhance.Sharpness(img2).enhance(1.25)
  img2.save("%s/temp/f%05d.png" % (frames_path, i+1))
  w, h = images[0].size
- cmd = 'ffmpeg -i %s/temp/f%%05d.png -c:v libx264 -pix_fmt yuv420p -vf scale=%d:%d %s'%(frames_path, w-(w%2), h-(h%2), output_path)
+ cmd = 'ffmpeg -y -i "%s/temp/f%%05d.png" -c:v libx264 -pix_fmt yuv420p -vf scale=%d:%d "%s"'%(frames_path, w-(w%2), h-(h%2), output_path)
  log('creating movie at %s' % output_path)
- os.system('rm %s' % output_path)
+ os.system('rm "%s"' % output_path)
  os.system(cmd)
- os.system('rm -rf %s/temp' % frames_path)
+ os.system('rm -rf "%s/temp"' % frames_path)
 
 
 def generate_video(frames_path, output_path, sat=1.0, con=1.0, sharp=1.0, match_hist=False, bitrate=None, cumulative=False, erase_frames=True):
- numframes = len([f for f in listdir(frames_path) if isfile(join(frames_path, f)) and f[-4:]=='.png'])
+ numframes = len(glob.glob('%s/f*.png'%frames_path))
  if numframes == 0:
  #warn("No frames found in %s"%frames_path)
  print("No frames found in %s"%frames_path)
  return
  log('creating %d-frame movie: %s -> %s'%(numframes, frames_path, output_path))
  if match_hist:
  avg_hist = get_average_histogram(frames_path)
- os.system('mkdir %s/temp' % frames_path)
+ os.system('mkdir "%s/temp"' % frames_path)
  for i in tqdm(range(numframes)):
  edited_filepath = "%s/temp/f%05d.png" % (frames_path, i+1)
  if not os.path.isfile(edited_filepath) or not cumulative:
@@ -61,16 +62,16 @@ def generate_video(frames_path, output_path, sat=1.0, con=1.0, sharp=1.0, match_
  img.save(edited_filepath)
  w, h = img.size
  wx, wy = w-(w%2), h-(h%2)
- ffmpeg_str = 'ffmpeg -i %s/temp/f%%05d.png -c:v libx264 -pix_fmt yuv420p -vf scale=%d:%d ' % (frames_path, wx, wy)
+ ffmpeg_str = 'ffmpeg -y -i "%s/temp/f%%05d.png" -c:v libx264 -pix_fmt yuv420p -vf scale=%d:%d ' % (frames_path, wx, wy)
  if bitrate is None:
- cmd = '%s %s'%(ffmpeg_str, output_path)
+ cmd = '%s "%s"'%(ffmpeg_str, output_path)
  else:
- cmd = '%s -b %d %s'%(ffmpeg_str, bitrate, output_path)
+ cmd = '%s -b %d "%s"'%(ffmpeg_str, bitrate, output_path)
  os.system(cmd)
  #if erase_frames:
  # os.system('rm %s' % output_path)
  if not cumulative:
- os.system('rm -rf %s/temp' % frames_path)
+ os.system('rm -rf "%s/temp"' % frames_path)
  log('Done making %s' % output_path)