initial commit

slitscan.py

@@ -0,0 +1,110 @@
+import argparse
+import cv2 as cv
+from itertools import count
+import numpy as np
+
+def duration(cap):
+    """
+    return the duration of the capture in seconds
+    """
+    nframes = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
+    fps = cap.get(cv.CAP_PROP_FPS)
+
+    return nframes / fps
+
+
+def capture_shape(cap):
+    """
+    return the shape of the video capture
+    """
+    nframes = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
+    height = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
+    width = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
+
+    return (nframes, height, width)
+
+
+def transposed_shape(shape, axis):
+    """
+    return the shape of the video capture after transposing the given axis to
+    become the new time axis
+    """
+    nframes, height, width = shape
+    if axis == 1:
+        new_shape = (height, nframes, width)
+    elif axis == 2:
+        new_shape = (width, height, nframes)
+
+    return new_shape
+
+
+def transpose_frame(infile, n, axis):
+    """
+    compute the nth frame of the input file after transposing the given axis to
+    become the new time axis
+    """
+    cap = cv.VideoCapture(infile)
+
+    shape = capture_shape(cap)
+    _, height, width = transposed_shape(shape, axis)
+    tframe = np.empty((height, width, 3), np.dtype('uint8'))
+
+    for i in count(0):
+        # NOTE: reading each frame of the input to produce one output frame is
+        # CPU-intensive, but to read the frames fewer times, we would have to
+        # store the frames in memory, which is memory-intensive
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        if axis == 1:
+            tframe[i] = frame[n]
+        elif axis == 2:
+            tframe[:, i] = frame[:, n]
+
+    cap.release()
+    return tframe
+
+
+def transpose_capture(infile, outfile, axis):
+    """
+    transpose the input file along the given axis, storing the output in
+    outfile
+    """
+    cap = cv.VideoCapture(infile)
+
+    seconds = duration(cap)
+    shape = capture_shape(cap)
+    nframes, height, width = transposed_shape(shape, axis)
+    fps = nframes / seconds
+
+    cap.release()
+
+    fourcc = cv.VideoWriter.fourcc(*'h264')
+    out = cv.VideoWriter(outfile, fourcc, fps, (width, height))
+
+    for i in range(0, nframes):
+        tframe = transpose_frame(infile, i, axis)
+        out.write(tframe)
+
+    out.release()
+    cv.destroyAllWindows()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("source")
+    parser.add_argument("dest")
+
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("--height", action='store_true')
+    group.add_argument("--width", action='store_true')
+
+    args = parser.parse_args()
+
+    if args.height:
+        axis = 1
+    elif args.width:
+        axis = 2
+
+    transpose_capture(args.source, args.dest, axis)