code.delx.au - transcoding/blob - video-transform

   1 #!/usr/bin/python2
   2
   3 from __future__ import division
   4
   5 import os
   6 import re
   7 import shutil
   8 import subprocess
   9 import sys
  10
  11 TMP_DIR = None
  12 DEST_DIR = None
  13 SOURCE_DIR = None
  14 DRY_RUN = False # this will still delete caches
  15
  16 VIDEO_FPS = 25
  17 AUDIO_SAMPLE_RATE = 48000
  18 ASPECT_RATIO = "4/3"
  19
  20
  21 def mkarg(arg):
  22     if re.match("^[a-zA-Z0-9\-\\.,/@_:=]*$", arg):
  23         return arg
  24
  25     if "'" not in arg:
  26         return "'%s'" % arg
  27     out = "\""
  28     for c in arg:
  29         if c in "\\$\"`":
  30             out += "\\"
  31         out += c
  32     out += "\""
  33     return out
  34
  35 def fail(line_count, msg):
  36     raise Exception(msg + " on line %d" % line_count)
  37
  38 def convert_frame_to_sample(frame):
  39     return frame * AUDIO_SAMPLE_RATE / VIDEO_FPS
  40
  41 def run_cmd(cmd):
  42     print "$", " ".join(map(mkarg, cmd))
  43     if DRY_RUN:
  44         return
  45     print
  46     ret = subprocess.Popen(cmd).wait()
  47     if ret != 0:
  48         print >>sys.stderr, "Failed on command", cmd
  49         raise Exception("Command returned non-zero: " + str(ret))
  50
  51 def read_file_contents(filename):
  52     try:
  53         f = open(filename)
  54         data = f.read().strip()
  55         f.close()
  56         return data
  57     except IOError:
  58         return None
  59
  60 def explode_video_to_png(source):
  61     image_cache_desc = os.path.join(TMP_DIR, "image_cache.txt")
  62     image_cache_dir = os.path.join(TMP_DIR, "image_cache")
  63
  64     # Do nothing if the current cache is what we need
  65     current_cache = read_file_contents(image_cache_desc)
  66     if source == current_cache:
  67         return image_cache_dir
  68
  69     # Remove if necessary
  70     if os.path.exists(image_cache_dir):
  71 ###     print "Confirm removal of image cache:", current_cache
  72 ###     ok = raw_input("(Y/n) ")
  73 ###     if ok != "Y":
  74 ###         print "Exiting..."
  75 ###         sys.exit(2)
  76         shutil.rmtree(image_cache_dir)
  77
  78     cmd = [
  79         "mplayer",
  80         "-vo", "png:outdir=%s" % image_cache_dir,
  81         "-nosound",
  82         "-noconsolecontrols",
  83         "-noconfig", "user",
  84         "-benchmark",
  85         source,
  86     ]
  87     run_cmd(cmd)
  88
  89     # Cache has been created, save the description
  90     f = open(image_cache_desc, "w")
  91     f.write(source)
  92     f.close()
  93
  94     return image_cache_dir
  95
  96
  97 def explode_video_to_wav(source):
  98     audio_cache_desc = os.path.join(TMP_DIR, "audio_cache.txt")
  99     audio_cache_file = os.path.join(TMP_DIR, "audio_cache.wav")
 100
 101     # Do nothing if the current cache is what we need
 102     if source == read_file_contents(audio_cache_desc):
 103         return audio_cache_file
 104
 105     cmd = [
 106         "mencoder",
 107         "-oac", "pcm",
 108         "-ovc", "copy",
 109         "-of", "rawaudio",
 110         "-o", audio_cache_file + ".raw",
 111         source,
 112     ]
 113     run_cmd(cmd)
 114
 115     cmd = [
 116         "sox",
 117         "-r", str(AUDIO_SAMPLE_RATE), "-b", "16", "-e", "signed-integer",
 118         audio_cache_file + ".raw",
 119         audio_cache_file,
 120     ]
 121     run_cmd(cmd)
 122
 123     # Cache has been created, save the description
 124     f = open(audio_cache_desc, "w")
 125     f.write(source)
 126     f.close()
 127
 128     return audio_cache_file
 129
 130 def apply_audio_effects(source, dest, crop_start, crop_end, audio_normalize):
 131     cmd = [
 132         "sox",
 133         source,
 134         dest,
 135     ]
 136     if audio_normalize:
 137         cmd += ["gain", "-n"]
 138     if crop_start and crop_end:
 139         c = convert_frame_to_sample
 140         cmd += ["trim", "%ds" % c(crop_start), "%ds" % c(crop_end - crop_start)]
 141     run_cmd(cmd)
 142
 143 def apply_single_image_effects(source_file, dest_file, color_matrix):
 144     cmd = [
 145         "convert",
 146         source_file,
 147         "-color-matrix", color_matrix,
 148         dest_file,
 149     ]
 150     run_cmd(cmd)
 151
 152 def apply_image_effects(source_dir, crop_start, crop_end, color_matrix):
 153     dest_dir = os.path.join(TMP_DIR, "image_processed")
 154     if os.path.exists(dest_dir):
 155         shutil.rmtree(dest_dir)
 156     os.mkdir(dest_dir)
 157
 158     inframe = crop_start
 159     outframe = 0
 160     while inframe <= crop_end:
 161         source_file = os.path.join(source_dir, str(inframe+1).zfill(8) + ".png")
 162         dest_file = os.path.join(dest_dir, str(outframe+1).zfill(8) + ".png")
 163         if color_matrix:
 164             apply_single_image_effects(source_file, dest_file, color_matrix)
 165         else:
 166             os.link(source_file, dest_file)
 167         inframe += 1
 168         outframe += 1
 169
 170     return dest_dir
 171
 172 def combine_audio_video(audio_file, image_dir, dest):
 173     cmd = [
 174         "mencoder",
 175         "mf://%s/*.png" % image_dir,
 176         "-audiofile", audio_file,
 177         "-force-avi-aspect", ASPECT_RATIO,
 178         "-vf", "harddup",
 179         "-af", "channels=1",
 180         "-ovc", "lavc",
 181         "-lavcopts", "vcodec=ffv1:ilme:ildct",
 182         "-oac", "pcm",
 183         "-o", dest,
 184     ]
 185     run_cmd(cmd)
 186
 187
 188 class Job(object):
 189     def __init__(self):
 190         self.source = None
 191         self.dest = None
 192         self.crop_start = None
 193         self.crop_end = None
 194         self.color_matrix = None
 195         self.audio_normalize = True
 196
 197     def set_source(self, arg):
 198         self.source = os.path.join(SOURCE_DIR, arg)
 199
 200     def set_dest(self, arg):
 201         self.dest = os.path.join(DEST_DIR, arg)
 202         if not self.dest.endswith(".avi"):
 203             self.dest += ".avi"
 204
 205     def set_crop(self, arg):
 206         a, b = arg.split("-")
 207         self.crop_start = int(a)
 208         self.crop_end = int(b)
 209
 210     def set_colormatrix(self, arg):
 211         [float(x) for x in arg.split(" ") if x] # check it's valid
 212         self.color_matrix = arg
 213
 214     def set_whitecolor(self, arg):
 215         arg = arg.split(" ")
 216         color = arg[0]
 217         r = 0xff / int(color[0:2], 16)
 218         g = 0xff / int(color[2:4], 16)
 219         b = 0xff / int(color[4:6], 16)
 220         # don't change the brightness
 221         avg = (r + g + b) / 3
 222         if (avg - 1) > 0.02:
 223             diff = avg - 1.0
 224             r -= diff
 225             g -= diff
 226             b -= diff
 227         if len(arg) == 2:
 228             brightness = float(arg[1])
 229             r *= brightness
 230             g *= brightness
 231             b *= brightness
 232         self.set_colormatrix("%.3f 0 0  0 %.3f 0  0 0 %.3f" % (r, g, b))
 233
 234     def set_audionormalize(self, arg):
 235         self.audio_normalize = int(arg)
 236
 237     def validate(self, line_count, unique):
 238         if self.dest in unique:
 239             fail(line_count, "Non-unique output file: " + self.dest)
 240         if self.source is None:
 241             fail(line_count, "Missing source")
 242         if self.dest is None:
 243             fail(line_count, "Missing dest")
 244         if not os.path.isfile(self.source):
 245             fail(line_count, "Unable to find source: " + self.source)
 246
 247     def is_done(self):
 248         return os.path.isfile(self.dest)
 249
 250     def run(self):
 251         image_cache_dir = explode_video_to_png(self.source)
 252         image_dir = apply_image_effects(image_cache_dir, self.crop_start, self.crop_end, self.color_matrix)
 253
 254         audio_cache_file = explode_video_to_wav(self.source)
 255         audio_file = os.path.join(TMP_DIR, "audio_processed.wav")
 256         apply_audio_effects(audio_cache_file, audio_file, self.crop_start, self.crop_end, self.audio_normalize)
 257
 258         combine_audio_video(audio_file, image_dir, self.dest+".tmp")
 259         os.rename(self.dest+".tmp", self.dest)
 260
 261     def __str__(self):
 262         return "Job :: %s (%s)" % (self.dest, self.source)
 263
 264 def main(frames):
 265     jobs = []
 266     unique = set()
 267
 268     f = open(frames)
 269
 270     job = None
 271     count = 0
 272
 273     def append_job():
 274         if job is None:
 275             return
 276         job.validate(count, unique)
 277         if job.is_done():
 278             print "Skipping", job
 279         else:
 280             jobs.append(job)
 281
 282     for line in f:
 283         count += 1
 284         line = line.strip()
 285         if line.startswith("#"):
 286             continue
 287         if not line:
 288             if job is not None:
 289                 append_job()
 290                 job = None
 291             continue
 292
 293         if job is None:
 294             job = Job()
 295         cmd, arg = line.split(" ", 1)
 296         f = getattr(job, "set_"+cmd, None)
 297         if not f:
 298             fail(count, "Invalid command: " + cmd)
 299         try:
 300             f(arg)
 301         except Exception, e:
 302             fail(count, str(e))
 303
 304     # trailing job...
 305     append_job()
 306
 307     # optimise image and audio cache usage, use the current cache first if it exists
 308     current_image_cache = read_file_contents(os.path.join(TMP_DIR, "image_cache.txt"))
 309     jobs.sort(key=lambda job: (job.source if job.source != current_image_cache else "", job.dest))
 310     for job in jobs:
 311         print "\n\n\nStarted job:", job, "\n\n"
 312         job.run()
 313
 314 if __name__ == "__main__":
 315     try:
 316         frames = sys.argv[1]
 317         SOURCE_DIR = sys.argv[2]
 318         DEST_DIR = sys.argv[3]
 319         TMP_DIR = sys.argv[4]
 320     except IndexError:
 321         print >>sys.stderr, "Usage: %s frames.txt source_dir dest_dir tmp_dir" % sys.argv[0]
 322         sys.exit(1)
 323
 324     main(frames)
 325