timeslice.md

timeslice.py - a script for generating timeslices from videos

timeslice.py

#!/usr/bin/env python

import argparse
import json
import math
import os
import subprocess
import sys

import numpy as np

from moviepy.editor import *
from PIL import Image, ImageDraw, ImageFont, ImageStat
from PIL.ExifTags import TAGS as id_names

class TimeSlice(object):
    def parse_arguments(self):
        self.parser = argparse.ArgumentParser()
        self.parser.add_argument('--dry_run',
            action='store_true',
            help='Do calculations (do not output composite)', )
        self.parser.add_argument('--reverse',
            action='store_true',
            help='Assemble slides with the oldest on the right', )
        self.parser.add_argument('--verbose',
            action='store_true',
            help='Output more info while running', )
        self.parser.add_argument('--luminance',
            action='store_true',
            help='Determine slices by luminance, not time (slower)', )
        self.parser.add_argument('--vertical',
            action='store_true',
            help='Slice from top to bottom, not left to right' )
        self.parser.add_argument('--label',
            action='store_true',
            help='Overlay time labels at top of image', )

        group = self.parser.add_mutually_exclusive_group(required=True)
        group.add_argument('--slices',
            action='store',
            type=int,
            help='Number of slices', )
        group.add_argument('--size',
            action='store',
            type=int,
            help='Size of slices, px', )

        self.parser.add_argument('path',
            help="Path of a video or directory", )

        self.parser.parse_args(namespace=self)

        # initialise byproduct internal variables
        self.time = not self.luminance
        self.isdir = os.path.isdir(self.path)

        self.check_args()

    def check_args(self):
        if not (self.slices or self.size):
            print "Either slices or width argument is required"
            sys.exit()
        if (not self.isdir) and (self.label):
            print "Cannot add labels from video"
            sys.exit()

    def load_frames(self):
        if self.isdir:
            images = os.listdir(self.path)
            self.images = [image for image in images if image.lower().endswith('.jpg')]
            if not self.images:
                warn("No JPG images in {}".format(self.path))
            self.images = sorted(self.images)
        else:
            # TODO is video actually video?
            self.clip = VideoFileClip(self.path)
            if self.verbose:
                print "Got clip with duration %.2f s" % self.clip.duration

    def get_root(self):
        path = self.path
        if self.isdir and path.endswith('/'):
            path = path.rstrip('/')

        (self.dir, self.file) = os.path.split(path)
        root = self.file.rsplit('.', 1)[0]
        return root

    def set_up_composite(self):
        if self.isdir:
            if not self.images:
                self.load_frames()
            final = Image.open(os.path.join(self.path, self.images[-1]))
            self.composite = Image.new("RGBA", final.size)
        else:
            if not self.clip:
                self.load_frames()
            self.composite = Image.new("RGBA", self.clip.size)

    def process_args(self):
        # TODO handle offset / files (see timeslice.py, previous versions)
        offset = 0

        if not self.composite:
            self.set_up_composite()

        (self.w, self.h) = self.composite.size

        if self.verbose:
            print "Output size {}x{}".format(self.w, self.h)

        if not self.vertical:
            size = self.w
        else:
            size = self.h

        if self.slices:
            self.slice_count = self.slices
            self.slice_size = size/self.slices
        if self.size:
            self.slice_count = size/self.size
            self.slice_size = self.size

        # TODO test for integer slice count and size

        # TODO warn if slices/width non-integer

        if self.label:
            self.font = self.init_font(max_width=self.slice_size*.9)

        if self.reverse:
            self.slice_size = -self.slice_size
        if self.verbose:
            print "Want %s slices" % self.slice_count
            print "Slice size %s px" % self.slice_size

    def init_font(self, font_name="OCRB.otf", max_width=None, max_chars=8):
        # this should possibly be looking elsewhere too
        font_path = os.path.join(
            os.path.expanduser("~"),
            'Library/Fonts',
            font_name,
        )

        # we need to calculate font_size if there's a max_width argument
        if max_width:
            font_size = 1
        else:
            font_size = 24

        self.font = ImageFont.truetype(font_path, font_size)
        if max_width:
            while self.font.getsize("x"*max_chars)[0]*2 < max_width:
                font_size += 1
                self.font = ImageFont.truetype(font_path, font_size)

            # reset to last size below criteria
            font_size -= 1
            self.font = ImageFont.truetype(font_path, font_size)

        self.font_size = font_size

        # get pixel size
        (self.font_width, self.font_height) = self.font.getsize("x"*max_chars)

        return self.font

    def determine_frames(self):
        if self.time:
            # TODO also allow width arg here
            if self.isdir:
                jump = len(self.images)/self.slice_count
                self.frames = [x*jump for x in range(0, self.slice_count)]
            else:
                jump = self.clip.duration*self.clip.fps/self.slice_count
                self.frames = [x*jump for x in range(0, self.slice_count)]
                if self.verbose:
                    print "Got frame indices %r" % self.frames

        if self.luminance:
            luminances = self.get_luminances()

            if self.verbose:
                print "Need to go from %s to %s" % (luminances[0], luminances[-1])

            diff = luminances[-1]-luminances[0]
            step = diff/self.slice_count

            if self.verbose:
                print "  Luminance step is %s" % (step)

            self.frames = [0]
            next = luminances[0]+step
            if self.verbose:
                print "  Using first, luminance %s target %s" % (luminances[0], next)

            for frame, luminance in enumerate(luminances):
                use_frame = False
                if step > 0: # we want a darker frame; is it?
                    if next <= luminance:
                        use_frame = True
                else:
                    if next >= luminance:
                        use_frame = True

                if use_frame:
                    next = next+step
                    if self.verbose:
                        if self.isdir:
                            print "  Using frame %s luminance %s new target %s" % (frame, luminance, next)
                        else:
                            time = float(frame)/self.clip.fps
                            print "  Using frame %s (time %s) luminance %s new target %s" % (frame, time, luminance, next)
                    self.frames.append(frame)

    def get_luminances(self):
        root = self.get_root()
        json_path = os.path.join(self.dir, "%s.json" % root)

        if os.path.exists(json_path):
            print "Loading cached luminances"
            f = open(json_path, 'r')
            luminances = json.load(f)
            f.close()
            return luminances

        if self.verbose:
            print "Calculating luminances"

        luminances = []
        if self.isdir:
            for image in self.images:
                print "  Loading {}".format(image)
                image = Image.open(os.path.join(self.path, image))
                luminances.append(self.find_brightness_image(image))
        else:
            for frame in self.clip.iter_frames():
                luminances.append(self.find_brightness_frame(frame))

        if self.verbose:
            print "Caching luminances to %s" % json_path
        f = open(json_path, 'w')
        json.dump(luminances, f)
        f.close()

        return luminances

    def find_brightness_image(self, image):
        stat = ImageStat.Stat(image)
        r, g, b = stat.mean
        # TODO tune factors, make constants?
        # saw 0.241, 0.691, 0.068 as factors elsewhere

        # the multiplication by the dimensions matches values from
        # numpy's sum.sum operations
        luminance = self.w*self.h* math.sqrt(0.299*(r**2) + 0.587*(g**2) + 0.114*(b**2))
        print "    Got brightness {}".format(luminance)
        return int(luminance)

    def find_brightness_frame(self, frame):
        rgb_total = np.sum(np.sum(frame, axis=0), axis=0)
        luminance = np.sqrt(0.299*np.square(rgb_total[0]) +
                            0.587*np.square(rgb_total[1]) +
                            0.114*np.square(rgb_total[2]))
        print "    Got brightness {}".format(luminance)
        return int(luminance)

    def make_composite(self):
        if self.verbose:
            print "Assembling composite"
        for idx, frame_index in enumerate(self.frames):
            if self.isdir:
                image_name =  self.images[frame_index]
                if self.verbose:
                    print "  Loading image {}, index {}, filename {}".format(
                        idx, frame_index, image_name,
                    )
                image = Image.open(os.path.join(self.path, image_name))
            else:
                time = float(frame_index)/self.clip.fps
                frame = self.clip.to_ImageClip(t=time)
                if self.verbose:
                    print "  Loading image {} frame {}, time {}".format(
                        idx, frame_index, time,
                    )
                image = Image.fromarray(np.uint8(frame.img))

            # Is there a better way to do this?
            if self.vertical:
                if not self.reverse:
                    # top to bottom
                    coords = (0, idx*self.slice_size,
                              self.w, self.slice_size+(idx*self.slice_size))
                else:
                    # bottom to top
                    coords = (0, self.h+self.slice_size+(idx*self.slice_size),
                              self.w, self.h+(idx*self.slice_size))
            else:
                if not self.reverse:
                    # left to right
                    coords = (idx*self.slice_size, 0,
                              self.slice_size+(idx*self.slice_size), self.h)
                else:
                    # right to left
                    coords = (self.w+self.slice_size+(idx*self.slice_size), 0,
                              self.w+(idx*self.slice_size), self.h)

            if self.verbose:
                if self.isdir:
                    print "  Pasting slice %s file %s to (%s)" % (idx, image_name, ", ".join([str(c) for c in coords]))
                else:
                    print "  Pasting slice %s at %.2f s to (%s)" % (idx, time, ", ".join([str(c) for c in coords]))

            slice = image.crop(coords)
            self.composite.paste(slice, coords)

            if self.verbose:
                print "  ... done pasting slice"
            if self.label and self.isdir:
                self.add_time_label(image, coords)

    def add_time_label(self, image, coords):
        if self.verbose:
            print "  * Adding label"
        time = self.get_time_from_image(image)
        if self.verbose:
            print "      Got time {}".format(time)
        overlay = self.get_overlay(time, coords)
        if self.verbose:
            print "      Got overlay"
        self.composite = Image.\
            alpha_composite(self.composite, overlay).\
            convert('RGBA')

    def get_time_from_image(self, image):
        """ Assumes the image is from PIL """
        raw_exif = image._getexif()
        tags = {name: raw_exif[id]
                for id, name in id_names.items()
                if id in raw_exif}

        date_tag = tags['DateTimeDigitized']
        time = str(date_tag).split(' ')[1]
        # date = str(date_tag).replace(':', '')

        return time

    def get_overlay(self, time, coords):
        # figure out the size of the text
        (font_width, font_height) = self.font.getsize(time)

        # make a blank image for the text, initialized to transparent text color
        overlay = Image.new('RGBA', self.composite.size, (255,255,255,0))
        draw = ImageDraw.Draw(overlay)

        # figure out where to put the text
        # TODO this does not support `--vertical`
        # TODO top/bottom argument?

        # top of image:
        text_top = font_height*.5
        # bottom of image:
        # text_top = self.composite.size[1]-(font_height*1.5)

        # abs() is needed for --reverse
        left_offset = (abs(self.slice_size)-font_width)/2.0
        text_left = coords[0]+left_offset

        # and the backing box
        box_top = text_top - (font_height*.2)
        box_left = text_left - (font_height*.2)
        box_bottom = box_top + font_height*1.4
        box_right = box_left + font_width + font_height*.4

        # add said backing box
        draw.rectangle(
            ((box_left, box_top), (box_right, box_bottom)),
            fill=(51,51,51,192),
        )

        # and draw then return
        white = (255,255,255,192)
        draw.text((text_left, text_top), time, white, font=self.font)
        if self.verbose:
            print "    Pasting text at {} {}".format(text_top, text_left)
            print "      over box at {} {} - {} {}".format(box_top,
                box_left, box_bottom, box_right)
        return overlay

    def save_composite(self):
        root = self.get_root()

        append = ""
        if self.luminance:
            append += "-l"
        if self.time:
            append += "-t"
        if self.reverse:
            append += "-r"
        if self.vertical:
            append += "-v"
        if self.label:
            # l for label and t for text taken, so, um
            append += "-z"

        print("Have self.dir {} and root {}".format(self.dir, root))

        filename = os.path.join(self.dir, "%s-%s%s.jpg" % (root, self.slice_count, append))
        self.composite.convert('RGB').save(filename)
        print "Saved %s" % filename
        subprocess.call(['open', filename])


if __name__ == '__main__':
    ts = TimeSlice()
    ts.parse_arguments()
    ts.load_frames()
    ts.set_up_composite()
    ts.process_args()
    ts.determine_frames()
    if not ts.dry_run:
        ts.make_composite()
        ts.save_composite()