added vector motion moshing with python

README.md

@@ -2,7 +2,8 @@ Datamoshing made easy.
 
 ## Requirements
 `mosh.py` requires `ffmpeg` to be installed.
-`vector_motion.py` and `style_transfer.py` rely on `ffedit` and `ffgac`, which can be downloaded from [ffglitch.org](https://ffglitch.org/)
+`vector_motion.py` and `style_transfer.py` rely on `ffedit` and `ffgac`, which can be downloaded from
+[ffglitch.org](https://ffglitch.org/)
 
 # Effects
 
@@ -10,26 +11,48 @@ Datamoshing made easy.
 This type of glitch creates the transition effect. Example:
 
     $ python mosh.py input.mp4 -s 40 -e 90 -o output.mp4
-removes all the i-frames from the input video starting at frame 40 and ending at frame 90, and outputs the final result to `output.mp4`
+removes all the i-frames from the input video starting at frame 40 and ending at frame 90, and outputs the final result
+to `output.mp4`
 
 ## Delta frame duplication
-Repeats a series of p-frames (aka delta frames), which can give a 'melting' effect. This type of glitch is triggered by the `-d` flag. Example:
+Repeats a series of p-frames (aka delta frames), which can give a 'melting' effect. This type of glitch is triggered by
+the `-d` flag. Example:
 
     $ python mosh.py input.mp4 -d 30 -s 30 -e 120 -o output.mp4
 
 copies frames with indexes [30, 59] and overwrites frames [60, 89] and [90, 119] with the copied data.
 
 ## Vector motion
-While the previous effects copy and delete whole frames, this one changes the actual frame data. As explained in [this article on ffglitch.org](https://ffglitch.org/2020/07/mv.html), you need to write a custom JavaScript file that can change the frame data. `vector_motion.py` is just a wrapper for `ffedit` and `ffgac` and makes moshing possible through only one command.
+While the previous effects copy and delete whole frames, this one changes the actual frame data. As explained in
+[this article on ffglitch.org](https://ffglitch.org/2020/07/mv.html), you need to write a custom JavaScript file
+that can change the frame data. `vector_motion.py` is just a wrapper for `ffedit` and `ffgac` and makes moshing
+possible through only one command.
 Example:
 
     $ python vector_motion.py input.mp4 -s your_script.js -o output.mp4
 
-I am planning to add support for moshing using python scripts in the near future (but this will be done internally through a js script).
+
+### Now also with Python!
+
+If you prefer to use python to glitch the frames instead, you can also specify a python script for the `-s` option.
+The script must contain a function called `mosh_frames` that takes as argument an array of frames (warning: some of the frames
+might be empty), where each non-empty frame represents a 3D array of shape (height, width, 2). The function should
+return an array containing the modified vectors. 
+
+`horizontal_motion_example.py` contains the equivalent python script of the js script presented in the
+[ffglitch tutorial](https://ffglitch.org/2020/07/mv.html).
+
+`average_motion_example.py` is the equivalent of this [ffglitch average motion tutorial](https://ffglitch.org/2020/07/mv_avg.html)
+using numpy. Neat!
+
 
 ## Style transfer
-I call style transfer combining the motion vectors of two videos (by simply adding them together). For example, applying the vector motion data of a person talking to a video of clouds can make it look as though the clouds are talking. 
-This script can also extract motion vector data from a video and write it to a file, or read motion data from file and apply it to a video.
+I call style transfer combining the motion vectors of two videos (by simply adding them together). For example,
+applying the vector motion data of a person talking to a video of clouds can make it look as though the clouds
+are talking. 
+
+This script can also extract motion vector data from a video and write it to a file, or read motion data from file and
+apply it to a video.
 
 Examples:
 

average_motion_example.py

@@ -0,0 +1,14 @@
+import numpy as np
+
+
+average_length = 10
+
+
+def average(frames):
+    if not frames:
+        return []
+    return np.mean(np.array([x for x in frames if x != []]), axis=0).tolist()
+
+
+def mosh_frames(frames):
+    return [average(frames[i + 1 - average_length: i + 1]) for i in range(len(frames))]

horizontal_motion_example.py

@@ -0,0 +1,11 @@
+def mosh_frames(frames):
+    for frame in frames:
+        if not frame:
+            continue
+
+        for row in frame:
+            for col in row:
+                # col contains the horizontal and vertical components of the vector
+                col[0] = 0
+
+    return frames

vector_motion.py

@@ -3,22 +3,49 @@
 import argparse
 import subprocess
 import os
+import importlib.util as imp
+from vector_util import get_vectors, apply_vectors
 
-parser = argparse.ArgumentParser()
-parser.add_argument('input_video', type=str, help='specifies input file')
-parser.add_argument('-s', type=str, dest='script_path', help='path to the js script', required=True)
-parser.add_argument('-g', type=str, default=1000, dest='gop_period', help='I-frame period (in frames)')
-parser.add_argument('-o', default='moshed.mpg', type=str, dest='output_video', help='output file for the moshed video')
-args = parser.parse_args().__dict__
 
-input_video = args['input_video']
-gop_period = args['gop_period']
-output_video = args['output_video']
-script_path = args['script_path']
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('input_video', type=str, help='specifies input file')
+    parser.add_argument('-s', type=str, dest='script_path', help='path to the script', required=True)
+    parser.add_argument('-g', type=str, default=1000, dest='gop_period', help='I-frame period (in frames)')
+    parser.add_argument('-o', default='moshed.mpg', type=str, dest='output_video',
+                        help='output file for the moshed video')
+    return parser.parse_args().__dict__
 
-subprocess.call(f'ffgac -i {input_video} -an -mpv_flags +nopimb+forcemv -qscale:v 0 -g {gop_period}' +
-                ' -vcodec mpeg2video -f rawvideo -y tmp.mpg', shell=True)
 
-subprocess.call(f'ffedit -i tmp.mpg -f mv -s {script_path} -o {output_video}', shell=True)
+def get_moshing_function(path):
+    spec = imp.spec_from_file_location('mod', path)
+    mod = imp.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod.mosh_frames
 
-os.remove('tmp.mpg')
+
+if __name__ == '__main__':
+    args = parse_args()
+    input_video = args['input_video']
+    gop_period = args['gop_period']
+    output_video = args['output_video']
+    script_path = args['script_path']
+    script_type = script_path[-2:]
+
+    if script_type == 'py':
+        try:
+            # import the function responsible for moshing from the specified script
+            func = get_moshing_function(script_path)
+
+            # method = '' specifies to just copy the new vectors over to the original video
+            apply_vectors(func(get_vectors(input_video)), input_video, output_video, method='')
+        except Exception as e:
+            # TODO: proper error handling
+            print(f'couldn\'t apply function mosh_frames from script: {script_path}:\n{e}')
+            exit(0)
+    else:
+        subprocess.call(f'ffgac -i {input_video} -an -mpv_flags +nopimb+forcemv -qscale:v 0 -g {gop_period}' +
+                        ' -vcodec mpeg2video -f rawvideo -y tmp.mpg', shell=True)
+
+        subprocess.call(f'ffedit -i tmp.mpg -f mv -s {script_path} -o {output_video}', shell=True)
+        os.remove('tmp.mpg')