shotorg.py

#!/usr/bin/env python

# shotorg 1.0
# written by Peter Delgobbo (@prdelgobbo)
# last updated May 2, 2014

# This program will use data generated from shotdetect to rip frames from a video
# file and organize them in their own directories, one for each shot. You can also
# make two types of visualizations for each shot: a one-row film strip-style montage
# or animated gifs (just for fun).

# Dependencies:
# shotdetect, ffmpeg, ImageMagick, BeautifulSoup

from bs4 import BeautifulSoup
import subprocess, os, os.path

### Functions

# Collects data for use in other functions
# The least elegant of all possible solutions, but it works
def get_data():
	home = raw_input("Input folder that contains video file: ")
	
	# ignore this, it's just to help me debug faster
	if home == "":
		home = "/home/prd/dh/uca"
		result = home + "/data/result.xml"
		soup = BeautifulSoup(open(result))
		video = home + "/uca.mp4"
		return {"home":home, "video":video, "result":result, "soup":soup}

	video = raw_input("Input name of video file: ")
	
	# change this to you shotdetect data directory
	result = home + "/data/result.xml"
	
	soup = BeautifulSoup(open(result))
	video = home + "/" + video

	print ""
	print "Working directory: " + home
	print "Video: " + video
	print "Shotdetect data file: " + result
	print ""

	# return a dictionary that we then feed to the other functions
	# home: base directory where the video file is
	# video: path to video
	# result: path to shotdetect xml output, result.xml
	# soup: result.xml read by BeautifulSoup
	return {"home":home, "video":video, "result":result, "soup":soup}

# Rips frames from a video file and arranges them into folders based on shot
def organize(data):
	# get some more user input that we will feed into ffmpeg
	res = raw_input("Input resolution for ripped frames (leave blank for default): ")
	fr = raw_input("Input sample rate (fps): ")

	# if no res is entered, read result.xml for default resolution
	if res == "":
		x = str(data["soup"].video).split()
		res = x[5]	

	for i in data["soup"].find_all("shot"):
		j = int(i.get("id"))
		start = str(float(i.get("msbegin")) / 1000.0) # convert ms to s
		dur = str(float(i.get("msduration")) / 1000.0)
		path = data["home"] + "/" + "%05d" % j # default 5 leading 0s (00000.png)
		
		subprocess.call(["mkdir", path]) # each shot gets its own directory

		# see ffmpeg documentation for more info:
		# https://www.ffmpeg.org/ffmpeg.html
		subprocess.check_call(["ffmpeg", "-ss", start, "-t", dur, "-i", data["video"], "-r", fr, "-s", res, "-f", "image2", "%05d.png"], cwd = path)

# Creates a 1-row montage of each shot
def montage(data):
	mpath = str(data["home"]) + "/montage" # output directory

	subprocess.call(["mkdir", mpath])

	for i in data["soup"].find_all("shot"):
		j = int(i.get("id"))
		path = data["home"] + "/%05d" % j
		out = mpath + "/" + "%05d.png" % j

		# skip any directory with 0 files
		if len([i for i in os.listdir(path) if os.path.isfile(path + "/" + i)]) > 0:
			# see ImageMagick documentation for more info:
			# http://www.imagemagick.org/Usage/montage/
			subprocess.check_call(["montage", "*.png", "-tile", "x1", "-geometry", "+0+0", out], cwd = path)

		print "Montage for shot " + str(j) + " completed"

# Create an animated gif from each shot
def loop():
	lpath = str(data["home"]) + "/loops" # output directory

	subprocess.call(["mkdir", lpath])

	for i in data["soup"].find_all("shot"):
		j = int(i.get("id"))
		path = data["home"] + "/%05d" % j
		out = lpath + "/" + "%05d.gif" % j

		# see ImageMagick documentation for more info:
		# http://www.imagemagick.org/script/convert.php
		subprocess.check_call(["convert", "-delay", "15", "-loop", "0", "*.png", out], cwd = path)

		print "Finished gif " + i.get("id")

# Mostly for debugging
def list_shots(data):
	for i in data["soup"].find_all("shot"):
		start = str(float(i.get("msbegin")) / 1000.0) 
		dur = str(float(i.get("msduration")) / 1000.0)

		print "Shot " + i.get("id") + " starts at " + start + " and lasts " + dur + " seconds."

### Script

x = False

while x == False:
	data = get_data()

	a = raw_input("Is this correct? [y/n] " )
	print ""

	if a == "y" or a == "Y" or a == "":
		x = True
		y = False

while y == False:
	b = raw_input("Organize, montage, loop, or exit? [o/m/l/x] ")

	if b == "o" or b == "O":
		organize(data)
	elif b == "m" or b == "M":
		montage(data)
	elif b == "l" or b == "L":
		loop(data)
	elif b == "x" or b == "X":
		y == True