temporally median filter videos

a.jpg

b.jpg

c.jpg

medianvideo.py

#
# medianvideo.py: temporally median filter videos
# Author: Bhautik Joshi ([email protected])
# Date: 16/01/2017
# License: MIT License
#
# Usage:
#  > python medianvideo.py <inputmovie.mp4> <windowsize>
#     inputmovie is your input movie, duh - eg. test.mov
#     windowsize is the number of frames you want to median filter over
# Outputs to: test.windowsize.mp4, defaults to 24fps
# If your frame rate is 24fps, then a window size of 120 means that
# you have a moving window of 5 seconds over which the median frame
# is calculated
#

import os, sys, shutil, glob, datetime
import numpy as np
import cv2

def getBase(fn):
  return (os.path.splitext(fn)[0])

def makeFreshDir(dirpath):
  if os.path.isdir(dirpath):
    shutil.rmtree(dirpath, ignore_errors=True)
  os.mkdir(dirpath)

def extractFrames(videoname, basename, outDir):
  outbase = os.path.join(outDir,basename)
  cmd = 'ffmpeg -i '+videoname+' -qscale:v 2 '+ outbase+'_%04d.jpg'
  print cmd
  os.system(cmd)

def loadImageAsArray(infilename):
  im = cv2.imread(infilename)
  return im

def saveImageFromArray(npdata, outfilename):
  cv2.imwrite(outfilename, npdata)

def createVideo(outDir, base, windowSize):
  outBase = os.path.join(outDir, base + '%04d.jpg')
  outmovie = base + '.' + str(windowSize) + '.mp4'
  # change the param afer -r from 24 to desired framerate if you want something different
  cmd = 'ffmpeg -start_number '+ str(windowSize+1)+' -i '+outBase+' -vcodec mpeg4 -r 24 ' + outmovie
  os.system(cmd)

def medianFrames(inDir, windowSize, outDir, base):
  fileList = glob.glob(os.path.join(inDir, '*.jpg'))
  counter = 0

  image_arrays = []
  for i in range(0, len(fileList)-windowSize + 1):
    start = datetime.datetime.now()
    imageList = fileList[i:i+windowSize]
    if i==0:
      for image in imageList:
        image_array = loadImageAsArray(image)
        image_arrays.append(image_array)
    else:
      #[1,2,3]->[2,3,4]
      image_arrays.pop(0)
      image_array = loadImageAsArray(imageList[-1])
      image_arrays.append(image_array)

    # create image stack, median filter it
    image_stack = np.concatenate([im[..., None] for im in image_arrays], axis=3)
    median_image = np.median(image_stack, axis=3)
    outFile = os.path.join(outDir, base+str(i+windowSize+1).zfill(4)+".jpg")
    saveImageFromArray(median_image,outFile)
    end = datetime.datetime.now()
    delta = end-start
    remainingFrames = len(fileList)-windowSize + 1 - i
    delta *=  remainingFrames
    print outFile, 'remaining: ' + str(delta)


  createVideo(outDir,base,windowSize)

def main2():
  image_arrays = [loadImageAsArray("a.jpg"),loadImageAsArray("b.jpg"), loadImageAsArray("c.jpg")]
  image_stack = np.concatenate([im[..., None] for im in image_arrays], axis=3)
  median_image = np.median(image_stack, axis=3)
  saveImageFromArray(median_image,"test.jpg")

def main():
  videoname = sys.argv[1]
  windowSize = int(sys.argv[2])
  basename = getBase(videoname)
  tmpBase = "tmp"
  if not os.path.isdir(tmpBase):
    os.mkdir(tmpBase)

  tmpFramesIn = os.path.join(tmpBase, basename + ".framesIn")
  tmpFramesOut = os.path.join(tmpBase, basename + ".framesOut")

  if not os.path.isdir(tmpFramesIn):
    makeFreshDir(tmpFramesIn)
    extractFrames(videoname, basename, tmpFramesIn)

  makeFreshDir(tmpFramesOut)

  medianFrames(tmpFramesIn, windowSize, tmpFramesOut, basename)

if __name__ == "__main__":
    main()
    #main2()