Explorar el Código

ckeck speed up

taruu hace 2 meses
padre
commit
01bb070e3c
Se han modificado 1 ficheros con 32 adiciones y 11 borrados
  1. 32 11
      main.py

+ 32 - 11
main.py

@@ -1,7 +1,8 @@
 import asyncio
+import io
 from concurrent.futures import ProcessPoolExecutor
 from itertools import groupby
-
+import numpy as np
 import bitstring
 from linuxpy.video.device import Device
 from Crypto.Hash import BLAKE2b
@@ -11,9 +12,25 @@ from loguru import logger
 logger.level("INFO")
 
 
-def chunks(lst, n):
-    for i in range(0, len(lst), n):
-        yield lst[i: i + n]
+def split_array_numpy(data):
+    data_np = np.frombuffer(data, dtype=np.uint32)
+    first = data_np[::2].tobytes()
+    second = data_np[1::2].tobytes()
+    return first, second
+
+
+def xor_arrays_numpy(arr1, arr2):
+    a = np.frombuffer(arr1, dtype=np.uint32)
+    b = np.frombuffer(arr2, dtype=np.uint32)
+    return np.bitwise_xor(a, b).tobytes()
+
+
+def chunks(in_bytes: bytes, take_n: int) -> bytes:
+    bytes_buffer = io.BytesIO(in_bytes)
+    bytes_send = bytes_buffer.read(take_n)
+    while len(bytes_send) == take_n:
+        yield bytes_send
+        bytes_send = bytes_buffer.read(take_n)
 
 
 def extract_lsbs(data):
@@ -22,6 +39,7 @@ def extract_lsbs(data):
     buffer = []
 
     data = [k for k, _ in groupby(data)]
+    data = bytes(data)
     data = chunks(data, 8)
 
     for chunk in data:
@@ -52,14 +70,16 @@ def whiten(data: bytes) -> bytes:
 
 def video2_extractor(frame_bytes: bytes):
     logger.debug(f"frames_in {len(frame_bytes)}")
-    chunked_bytes = chunks(frame_bytes, 8)
-    newdata = b""
+    # chunked_bytes = chunks(frame_bytes, 8)
+    # newdata = b""
 
-    for chunk in chunked_bytes:
-        newdata += bytes(a ^ b for a, b in zip(chunk[:4], chunk[4:]))
+    first_part, second_part = split_array_numpy(frame_bytes)
+    new_data = xor_arrays_numpy(first_part, second_part)
+    # for chunk in chunked_bytes:
+    #     newdata += bytes(a ^ b for a, b in zip(chunk[:4], chunk[4:]))
 
-    logger.debug(f"newdata {len(newdata)}")
-    out_bytes = extract_lsbs(newdata)
+    logger.debug(f"new data {len(new_data)}")
+    out_bytes = extract_lsbs(new_data)
     out_bytes = whiten(out_bytes)
 
     logger.info(f"Sample ready: {len(out_bytes)}b.")
@@ -87,7 +107,8 @@ class FrameTakeWorker:
         return next(self.frame_gen)
 
     def __del__(self):
-        self.frame_gen.close()
+        if self.frame_gen:
+            self.frame_gen.close()
 
     async def frame_extractor(self, frame_bytes: bytes):
         loop = asyncio.get_event_loop()