process.py

from __future__ import division
import matplotlib
matplotlib.use('Agg')

import time
import sys
import math
import glob
import os
import cv2

import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np

from itertools import cycle

cv2.ocl.setUseOpenCL(False)

def points(a,b):
    return np.array(np.meshgrid(a, b)).reshape(2,len(a)*len(b)).T

def similarity_metric(volume):
    fvolume = volume.astype(float)
    diffs = np.zeros((len(fvolume), len(fvolume)))
    print("Calculating similarity metric for volume of length {0}".format(len(fvolume)))

    for i in range(len(fvolume)):
        if (i % 10) == 0:
            print("Calculating row {0} of {1}".format(i, len(fvolume)))
        for j in range(len(fvolume)):
            if i != j:
                diffs[i,j] = ((fvolume[i] - fvolume[j])**2).sum() ** 0.5
            else:
                diffs[i,j] = 0.0

    diffs /= diffs.mean()

    #output = np.sqrt(((images - images[:, np.newaxis])**2).sum(axis=(2,3,4)))
    #output /= output.mean()
    return diffs

def transition_diff(ssd_difference):
    output = np.zeros((ssd_difference.shape[0] - 4,
                       ssd_difference.shape[1] - 4),
                      dtype=ssd_difference.dtype)
    kernel = np.eye(5) * binomial_filter()
    output[:] = cv2.filter2D(ssd_difference, -1, kernel)[2:-2,2:-2]
    return output

def biggest_loop(trans_diff, alpha):
    start = 0
    end = 0
    largest_score = 0
    for i in range(trans_diff.shape[0]):
        for j in range(trans_diff.shape[1]):
            diff = trans_diff[j,i]
            score = (1.0*alpha) * (j - i) - diff
            if score > largest_score:
                largest_score, start, end = score, i, j
    return start, end

def synthesize_loop(video_volume, start, end):
    print("Synthesizing volume with start {0} and end {1}".format(start, end))
    return video_volume[start:end+1]

def binomial_filter():
    return np.array([1 / 16., 1 / 4., 3 / 8., 1 / 4., 1 / 16.], dtype=float)

def get_frames(root, masks=False, excluded=None, green=False, sensitivity=33, method="MOG"):
    search = map(lambda x: "{0}/*.{1}".format(root, x), ["mp4", "mov", "m4v"])
    files = sum(map(glob.glob, search), [])
    if len(files) == 0:
        print("Could not find an input file for {0}".format(root))
        exit(1)
    if os.path.exists("{0}/expanded/0000.png".format(root)):
        print("Frames already created for {0}, skipping...".format(root))
        return
    excluded = excluded or list()
    dispatch = {
        'MOG': cv2.bgsegm.createBackgroundSubtractorMOG,
        'MOG2': cv2.createBackgroundSubtractorMOG2,
        'GMG': cv2.bgsegm.createBackgroundSubtractorGMG,
        'KNN': cv2.createBackgroundSubtractorKNN,
        'GREEN': createBackgroundSubtractorChromaKey,
    }

    cap = cv2.VideoCapture(files[0])
    fgbg = dispatch[method]()
    framedir = "{0}/expanded".format(root)
    maskdir = framedir.replace("expanded", "masks")

    if not os.path.exists(framedir): os.mkdir(framedir)
    if not os.path.exists(maskdir) and masks: os.mkdir(maskdir)

    count = 0
    ret, frame = cap.read()
    fgmask = fgbg.apply(frame)
    while cap.isOpened():
        ret, frame = cap.read()
        if ret != True: break

        if masks:
            if not count in excluded:
                fgmask = fgbg.apply(frame, fgmask)
            else:
                fgmask = np.zeros(frame.shape)

            if method == "GREEN":
                mask = fgmask.copy()
                mask = cv2.normalize(mask, mask, 0, 1, cv2.NORM_MINMAX)
                frame = frame * cv2.merge([mask, mask, mask])

            out = "%s/%04d.png" % (maskdir, count)
            imwrite(out, fgmask)
        out = "%s/%04d.png" % (framedir, count)
        imwrite(out, frame)
        count += 1
    cap.release()

class ChromaKey(object):
    def __init__(self):
        self.sensitivity = 33
    def apply(self, frame, fgmask=None):
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        lower = np.array([60-self.sensitivity, 100, 90])
        upper = np.array([60+self.sensitivity, 255, 255])
        mask = cv2.inRange(hsv, lower, upper)
        return cv2.bitwise_not(mask)

def createBackgroundSubtractorChromaKey():
    return ChromaKey()

def get_texture(images, alpha=1.0):
    print("calculating similarity matrix...")
    ssd_diff = similarity_metric(images)
    print("trans_diff...")
    trans_diff = transition_diff(ssd_diff)
    print("biggest_loop ...")
    start, end = biggest_loop(trans_diff, alpha)

    print("diff3...")
    diff3 = np.zeros(trans_diff.shape, float)

    for i in range(trans_diff.shape[0]):
        for j in range(trans_diff.shape[1]):
            diff3[i, j] = alpha * (i - j) - trans_diff[i, j]

    return (viz_difference(ssd_diff),
            viz_difference(trans_diff),
            viz_difference(diff3),
            synthesize_loop(images, start + 2, end + 2))

def viz_difference(diff):
    return (((diff - diff.min()) /
             (diff.max() - diff.min())) * 255).astype(np.uint8)

def frames(first, second, size):
    # gather all images
    f1 = glob.glob("{0}/expanded/*.png".format(first))
    f2 = glob.glob("{0}/expanded/*.png".format(second))
    masks = glob.glob("{0}/masks/*.png".format(second))

    # resolve mismatches in frames by cycling
    if len(f1) < len(f2):
        c = cycle(f1)
        f1 = [next(c) for _ in range(len(f2))]
    elif len(f2) < len(f1):
        c = cycle(f2)
        m = cycle(masks)
        f2 = [next(c) for _ in range(len(f1))]
        masks = [next(m) for _ in range(len(f1))]

    files = zip(f1, f2, masks)

    for f1, f2, m in files:
        black = cv2.resize(cv2.imread(f1), size)
        white = cv2.resize(cv2.imread(f2), size)
        mask  = cv2.resize(cv2.imread(m), size)
        #mask[mask > 128] = 255
        #mask[mask <= 128] = 0
        mask = mask.astype(float) / 255
        yield (black, white, mask)

def depth(shape):
    return int(math.floor(math.log(min(shape), 2))) - 4

def gauss(black, white, mask, depth):
    pyr_white, pyr_black, pyr_mask = [white], [black], [mask]

    for _ in range(depth):
        pyr_white.append(cv2.pyrDown(pyr_white[-1]))
        pyr_black.append(cv2.pyrDown(pyr_black[-1]))
        pyr_mask.append(cv2.pyrDown(pyr_mask[-1]))
    return pyr_white, pyr_black, pyr_mask

def laplacian(gblack, gwhite):
    lblack, lwhite = [gblack[-1]], [gwhite[-1]]

    for i in range(len(gblack)-1, 0, -1):
        size = gblack[i-1].shape[1], gblack[i-1].shape[0]
        GE = cv2.pyrUp(gblack[i], dstsize=size)
        L = cv2.subtract(gblack[i-1], GE)
        lblack.append(L)

        size = gwhite[i-1].shape[1], gwhite[i-1].shape[0]
        GE = cv2.pyrUp(gwhite[i], dstsize=size)
        L = cv2.subtract(gwhite[i-1], GE)
        lwhite.append(L)
    return lblack, lwhite

def blend(first, second, outdir, morph=False, size=(400,300)):
    print("Blending frames....")
    if morph:
        print("Morphing masks enabled")
    nframes = []
    mcount = 0
    for black, white, mask in frames(first, second, size=size):
        if morph:
            #kernel = np.ones((5,5), dtype=np.uint8)
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(3,3))
            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=3)
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=3)
            outfile = "{0}/morphed/{1:04d}.png".format(outdir, mcount)
            imwrite(outfile, mask)
            mcount += 1
        d = depth(white.shape[:-1])
        gwhite, gblack, gmask = gauss(white, black, mask, d)
        lwhite, lblack = laplacian(gwhite, gblack)

        pyr_blended = []
        gmask.reverse()
        for w, b, g in zip(lwhite, lblack, gmask):
            output = g * b + (1 - g) * w
            pyr_blended.append(output)

        img = pyr_blended[0]
        for i in range(1,len(pyr_blended)):
            size = pyr_blended[i].shape[1], pyr_blended[i].shape[0]
            img = cv2.pyrUp(img, dstsize=size)
            img = cv2.add(img, pyr_blended[i])
        nframes.append(img)
    return nframes

def normalize(frame):
    norm = cv2.normalize(frame,
                         frame,
                         alpha=0,
                         beta=255,
                         norm_type=cv2.NORM_MINMAX,
                         dtype=cv2.CV_8UC3)
    return norm

def create_texture(nframes, outdir, alpha=1.0):
    print("Saving textures to {0}".format(outdir))
    fourcc = cv2.VideoWriter_fourcc(*'avc1')
    outfile = '{0}/texture/output.mp4'.format(outdir)
    out = cv2.VideoWriter(outfile, fourcc, 10, nframes[0].shape[-2:-4:-1])

    diff1, diff2, diff3, output = get_texture(nframes, alpha)

    imwrite("{0}/diffs/diff1.png".format(outdir), diff1)
    imwrite("{0}/diffs/diff2.png".format(outdir), diff2)
    imwrite("{0}/diffs/diff3.png".format(outdir), diff3)

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_axis_off()
    ims = []

    for idx, frame in enumerate(output):
        out.write(frame)
        outfile = "{0}/texture/frames/{1:04d}.png".format(outdir, idx)
        imwrite(outfile, frame)
        b, g, r = cv2.split(frame)
        rgb = cv2.merge([r,g,b])
        ims.append([ax.imshow(rgb)])
    out.release()

    # http://www.nooganeer.com/his/projects/image-processing/making-a-gif-with-opencv-and-scikit-image-in-python/
    ani = animation.ArtistAnimation(fig, ims, interval=2, repeat_delay=0, blit=False)
    outfile = "{0}/texture/animation.gif".format(outdir)
    ani.save(outfile, writer='imagemagick')

def imwrite(outfile, img):
    if not cv2.imwrite(outfile, img):
        print("ERROR: Failed to write image to {0}".format(outfile))

def write_frames(nframes, output):
    print("Saving frames to cache")
    for idx, frame in enumerate(nframes):
        imwrite("{0}/frames/{1:04d}.png".format(output, idx), frame)

def parse_args(args):
    alpha = 1.0
    if len(args) > 0:
        alpha = float(args[0])
    print("Using alpha {0}".format(alpha))
    return alpha

def setup(first, second, method):
    output = "results/{0}_{1}_{2}_{3}".format(first, second, method, time.strftime("%Y%m%d%H%M%S"))
    for d in ["diffs", "texture/frames", "morphed", "frames"]:
        os.makedirs("{0}/{1}".format(output, d))

    excluded = sum([range(11), range(56,120)], [])
    return output, excluded

def run(first, second, output, method):
    print("Getting frames for {0}....".format(first))
    get_frames(first)
    print("Getting frames for {0}....".format(second))
    get_frames(second, masks=True,
               excluded=None,
               green=True,
               sensitivity=33,
               method=method
    )

    print("Computing blended frames....")
    normed = [normalize(frame) for frame in
                    blend(first, second, output, morph=True)]
    nframes = np.array(normed)
    print("Saving {0} blended frames...".format(len(nframes)))
    write_frames(nframes, output)
    return nframes

def clean(inputs):
    for d in inputs:
        print("Cleaning out files in {0}".format(d))
        for out in ["masks", "expanded"]:
            for image in glob.glob("{0}/{1}/*.png".format(d, out)):
                os.remove(image)

def main(first="outside", second="dragon", method="MOG"):
    alpha = parse_args(sys.argv[1:])
    output, excluded = setup(first, second, method)

    nframes = run(first, second, output, method)
    create_texture(nframes, output, alpha)

if __name__ == "__main__":
    #for second in ["dragon", "crowd", "dolphins", "bloodborne", "walkers2", "walkers"]:
    for second in ["walkers2", "walkers"]:
        for method in ['MOG', 'MOG2', 'GMG', 'KNN', 'GREEN']:
            msg = "Running trial for first {0}, second {1}, using method {2}"
            print(msg.format("outside", second, method))
            clean(["outside", second])
            main("outside", second, method)