Try multithreading

1 files changed, 81 insertions, 64 deletions

diff --git a/decoder.py b/decoder.py
index 8166382..86563ef 100644
--- a/decoder.py
+++ b/decoder.py
@@ -1,4 +1,6 @@
 import argparse
+from multiprocessing.pool import Pool
+from threading import Thread
 import time
 import cv2
 import numpy as np
@@ -23,6 +25,9 @@ rs_bytes = frame_bytes - (frame_bytes + 254) // 255 * int(args.level * 255) - 4
 
 rsc = RSCodec(int(args.level * 255))
 decoder = Decoder.with_defaults(args.size, rs_bytes)
+data = None
+decoded = 0
+start_time = 0
 
 
 def find_corner(A, f):
@@ -43,78 +48,90 @@ def find_corner(A, f):
     return np.average(np.where(mask), axis=1), np.average(A[mask], axis=0).astype(np.float64)
 
 
-if isinstance(args.input, str) and args.input.isdecimal():
-    args.input = int(args.input)
-cap = cv2.VideoCapture(args.input)
-data = None
-start_time = 0
-decoded = 0
-while data is None:
-    try:
-        ret, raw_frame = cap.read()
-        if not ret:
-            print("End of stream")
-            break
-        # raw_frame is a uint8 BE CAREFUL
-        cv2.imshow("", raw_frame)
-        cv2.waitKey(1)
-        raw_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGB)
-
-        # Find positions and colors of corners
-        X, Y = raw_frame.shape[:2]
-        cx, cy = X // 4, Y // 4
-        widx, wcol = find_corner(raw_frame[:cx, :cy], lambda B: np.sum(B, axis=2) - 2 * np.std(B, axis=2))
-        ridx, rcol = find_corner(raw_frame[:cx, Y - cy :], lambda B: B[:, :, 0] - B[:, :, 1] - B[:, :, 2])
-        ridx[1] += Y - cy
-        gidx, gcol = find_corner(raw_frame[X - cx :, :cy], lambda B: B[:, :, 1] - B[:, :, 2] - B[:, :, 0])
-        gidx[0] += X - cx
-        bidx, bcol = find_corner(raw_frame[X - cx :, Y - cy :], lambda B: B[:, :, 2] - B[:, :, 0] - B[:, :, 1])
-        bidx[0] += X - cx
-        bidx[1] += Y - cy
+def rsc_decode(x):
+    return rsc.decode(x)[0]
 
-        # Find basis of color space
-        origin = (rcol + gcol + bcol - wcol) / 2
-        rcol -= origin
-        gcol -= origin
-        bcol -= origin
-        F = 255 * np.linalg.inv(np.stack((rcol, gcol, bcol)).T)
 
-        cch = cheight / 2 - 1
-        ccw = cwidth / 2 - 1
-        M = cv2.getPerspectiveTransform(
-            np.float32([np.flip(widx), np.flip(ridx), np.flip(gidx), np.flip(bidx)]),
-            np.float32(
-                [
-                    [ccw, cch],
-                    [args.width - ccw - 1, cch],
-                    [ccw, args.height - cch - 1],
-                    [args.width - ccw - 1, args.height - cch - 1],
-                ]
-            ),
-        )
-        frame = cv2.warpPerspective(raw_frame, M, (args.width, args.height))
-        # Convert to new color space
-        frame = (np.squeeze(F @ (frame - origin)[..., np.newaxis]) >= 160).astype(np.uint8)
-        frame = np.packbits(
-            np.concatenate(
-                (
-                    frame[:cheight, cwidth : args.width - cwidth].flatten(),
-                    frame[cheight : args.height - cheight].flatten(),
-                    frame[args.height - cheight :, cwidth : args.width - cwidth].flatten(),
-                )
-            )
-        )[:frame_bytes]
-        reshape_len = frame_bytes // 255 * 255
-        frame[:reshape_len] = np.ravel(frame[:reshape_len].reshape(255, reshape_len // 255), "F")
-        data = decoder.decode(bytes(rsc.decode(bytearray(frame ^ frame_xor))[0][: args.psize]))
+def ecc_decode(frame):
+    global data
+    global decoded
+    global start_time
+    try:
+        rsc_frame = bytearray(frame ^ frame_xor)
+        with Pool(processes=1) as p:
+            rsc_data = p.apply(rsc_decode, (rsc_frame,))
+        data = decoder.decode(bytes(rsc_data[: args.psize]))
         decoded += 1
         if start_time == 0:
             start_time = time.time()
         print("Decoded frame")
-    except KeyboardInterrupt:
-        break
     except Exception as e:
         print(e)
+
+
+if isinstance(args.input, str) and args.input.isdecimal():
+    args.input = int(args.input)
+cap = cv2.VideoCapture(args.input)
+while data is None:
+    frame_start_time = time.time()
+    ret, raw_frame = cap.read()
+    if not ret:
+        print("End of stream")
+        break
+    # raw_frame is a uint8 BE CAREFUL
+    cv2.imshow("", raw_frame)
+    cv2.waitKey(1)
+    raw_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGB)
+
+    # Find positions and colors of corners
+    X, Y = raw_frame.shape[:2]
+    cx, cy = X // 4, Y // 4
+    widx, wcol = find_corner(raw_frame[:cx, :cy], lambda B: np.sum(B, axis=2) - 2 * np.std(B, axis=2))
+    ridx, rcol = find_corner(raw_frame[:cx, Y - cy :], lambda B: B[:, :, 0] - B[:, :, 1] - B[:, :, 2])
+    ridx[1] += Y - cy
+    gidx, gcol = find_corner(raw_frame[X - cx :, :cy], lambda B: B[:, :, 1] - B[:, :, 2] - B[:, :, 0])
+    gidx[0] += X - cx
+    bidx, bcol = find_corner(raw_frame[X - cx :, Y - cy :], lambda B: B[:, :, 2] - B[:, :, 0] - B[:, :, 1])
+    bidx[0] += X - cx
+    bidx[1] += Y - cy
+
+    # Find basis of color space
+    origin = (rcol + gcol + bcol - wcol) / 2
+    rcol -= origin
+    gcol -= origin
+    bcol -= origin
+    F = 255 * np.linalg.inv(np.stack((rcol, gcol, bcol)).T)
+
+    cch = cheight / 2 - 1
+    ccw = cwidth / 2 - 1
+    M = cv2.getPerspectiveTransform(
+        np.float32([np.flip(widx), np.flip(ridx), np.flip(gidx), np.flip(bidx)]),
+        np.float32(
+            [
+                [ccw, cch],
+                [args.width - ccw - 1, cch],
+                [ccw, args.height - cch - 1],
+                [args.width - ccw - 1, args.height - cch - 1],
+            ]
+        ),
+    )
+    frame = cv2.warpPerspective(raw_frame, M, (args.width, args.height))
+    # Convert to new color space
+    frame = (np.squeeze(F @ (frame - origin)[..., np.newaxis]) >= 160).astype(np.uint8)
+    frame = np.packbits(
+        np.concatenate(
+            (
+                frame[:cheight, cwidth : args.width - cwidth].flatten(),
+                frame[cheight : args.height - cheight].flatten(),
+                frame[args.height - cheight :, cwidth : args.width - cwidth].flatten(),
+            )
+        )
+    )[:frame_bytes]
+    reshape_len = frame_bytes // 255 * 255
+    frame[:reshape_len] = np.ravel(frame[:reshape_len].reshape(255, reshape_len // 255), "F")
+    # Run ECC decoding in separate thread
+    Thread(target=ecc_decode, args=(frame,)).start()
+    print(time.time() - frame_start_time)
 cap.release()
 print(decoded)
 with open(args.output, "wb") as f: