created separate file for vector utils

style_transfer.py

@@ -1,84 +1,12 @@
 #!/usr/bin/env python3
 
 import argparse
-import subprocess
-import os
 import json
-
-
-def get_vectors(input_video):
-    # extract video data using ffedit
-    subprocess.call(f'ffedit -i {input_video} -f mv:0 -e tmp.json', shell=True)
-
-    # read the data we extracted
-    f = open('tmp.json', 'r')
-    raw_data = json.load(f)
-    f.close()
-    os.remove('tmp.json')
-
-    # read frame information
-    frames = raw_data['streams'][0]['frames']
-
-    # get vectors from each frame
-    vectors = []
-    for frame in frames:
-        try:
-            vectors.append(frame['mv']['forward'])
-        except:
-            vectors.append([])
-
-    return vectors
-
-
-def apply_vectors(vectors, input_video, output_video, method='add'):
-    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)
-
-    # open js file and read its contents
-    to_add = '+' if method == 'add' else ''
-    script_contents = '''
-    var vectors = [];
-    var n_frames = 0;
-
-    function glitch_frame(frame) {
-        let fwd_mvs = frame["mv"]["forward"];
-        if (!fwd_mvs || !vectors[n_frames]) {
-            n_frames++;
-            return;
-        }
-
-        for ( let i = 0; i < fwd_mvs.length; i++ ) {
-            let row = fwd_mvs[i];
-            for ( let j = 0; j < row.length; j++ ) {
-                let mv = row[j];
-                try {
-                    mv[0] ''' + to_add + '''= vectors[n_frames][i][j][0];
-                    mv[1] ''' + to_add + '''= vectors[n_frames][i][j][1];
-                } catch {}
-            }
-        }
-
-        n_frames++;
-    }
-    '''
-
-    script_path = 'apply_vectors.js'
-
-    # open js file and write the code
-    with open(script_path, 'w') as f:
-        f.write(script_contents.replace('var vectors = [];', f'var vectors = {json.dumps(vectors)};'))
-
-    # apply the effect
-    subprocess.call(f'ffedit -i tmp.mpg -f mv -s {script_path} -o {output_video}', shell=True)
-
-    # remove temp files
-    os.remove('apply_vectors.js')
-    os.remove('tmp.mpg')
+from vector_util import *
 
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument('-g', type=str, default=1000, dest='gop_period', help='I-frame period (in frames)')
     parser.add_argument('-v', type=str, default='', dest='vector_file',
                         help='file containing vector data to transfer', required=False)
     parser.add_argument('-e', type=str, default='', dest='extract_from',
@@ -92,7 +20,6 @@ def parse_args():
 if __name__ == '__main__':
     # get args
     args = parse_args()
-    gop_period = args['gop_period']
     vector_file = args['vector_file']
     extract_from = args['extract_from']
     transfer_to = args['transfer_to']
@@ -107,10 +34,7 @@ def parse_args():
 
     if extract_from:
         # step 1a: extract vectors
-        subprocess.call(f'ffgac -i {extract_from} -an -mpv_flags +nopimb+forcemv -qscale:v 0 -g {gop_period}' +
-                        ' -vcodec mpeg2video -f rawvideo -y tmp.mpg', shell=True)
-        vectors = get_vectors('tmp.mpg')
-        os.remove('tmp.mpg')
+        vectors = get_vectors(extract_from)
 
         # if we only have to extract the vectors, write to file and exit
         if transfer_to == '':

vector_util.py

@@ -0,0 +1,76 @@
+import os
+import subprocess
+import json
+
+
+def get_vectors(input_video):
+    # extract video data using ffedit
+    subprocess.call(f'ffgac -i {input_video} -an -mpv_flags +nopimb+forcemv -qscale:v 0 -g 1000' +
+                    ' -vcodec mpeg2video -f rawvideo -y tmp.mpg', shell=True)
+    subprocess.call(f'ffedit -i tmp.mpg -f mv:0 -e tmp.json', shell=True)
+    os.remove('tmp.mpg')
+
+    # read the data we extracted
+    f = open('tmp.json', 'r')
+    raw_data = json.load(f)
+    f.close()
+    os.remove('tmp.json')
+
+    # read frame information
+    frames = raw_data['streams'][0]['frames']
+
+    # get vectors from each frame
+    vectors = []
+    for frame in frames:
+        try:
+            vectors.append(frame['mv']['forward'])
+        except:
+            vectors.append([])
+
+    return vectors
+
+
+def apply_vectors(vectors, input_video, output_video, method='add'):
+    subprocess.call(f'ffgac -i {input_video} -an -mpv_flags +nopimb+forcemv -qscale:v 0 -g 1000' +
+                    ' -vcodec mpeg2video -f rawvideo -y tmp.mpg', shell=True)
+
+    # open js file and read its contents
+    to_add = '+' if method == 'add' else ''
+    script_contents = '''
+    var vectors = [];
+    var n_frames = 0;
+
+    function glitch_frame(frame) {
+        let fwd_mvs = frame["mv"]["forward"];
+        if (!fwd_mvs || !vectors[n_frames]) {
+            n_frames++;
+            return;
+        }
+
+        for ( let i = 0; i < fwd_mvs.length; i++ ) {
+            let row = fwd_mvs[i];
+            for ( let j = 0; j < row.length; j++ ) {
+                let mv = row[j];
+                try {
+                    mv[0] ''' + to_add + '''= vectors[n_frames][i][j][0];
+                    mv[1] ''' + to_add + '''= vectors[n_frames][i][j][1];
+                } catch {}
+            }
+        }
+
+        n_frames++;
+    }
+    '''
+
+    script_path = 'apply_vectors.js'
+
+    # open js file and write the code
+    with open(script_path, 'w') as f:
+        f.write(script_contents.replace('var vectors = [];', f'var vectors = {json.dumps(vectors)};'))
+
+    # apply the effect
+    subprocess.call(f'ffedit -i tmp.mpg -f mv -s {script_path} -o {output_video}', shell=True)
+
+    # remove temp files
+    os.remove('apply_vectors.js')
+    os.remove('tmp.mpg')