Browse Source

first work test

Taruu 4 months ago
parent
commit
fcf1359d16
1 changed files with 65 additions and 8 deletions
  1. 65 8
      main.py

+ 65 - 8
main.py

@@ -1,13 +1,18 @@
+import argparse
 import asyncio
 import asyncio
 import io
 import io
+import random
+import time
 from concurrent.futures import ProcessPoolExecutor
 from concurrent.futures import ProcessPoolExecutor
 from itertools import groupby
 from itertools import groupby
 import numpy as np
 import numpy as np
 import bitstring
 import bitstring
 from linuxpy.video.device import Device
 from linuxpy.video.device import Device
 from Crypto.Hash import BLAKE2b
 from Crypto.Hash import BLAKE2b
+import aiohttp
 
 
 from loguru import logger
 from loguru import logger
+import math
 
 
 logger.level("INFO")
 logger.level("INFO")
 
 
@@ -89,8 +94,8 @@ def video2_extractor(frame_bytes: bytes):
     return out_bytes
     return out_bytes
 
 
 
 
-class FrameTakeWorker:
-    def __init__(self, device_id=0, processes=3):
+class GeneratorWorker:
+    def __init__(self, pool_url: str, auth: str, device_id=0, processes=3, frame_timeout=1, push_timeout=10):
         self.processes = processes
         self.processes = processes
         self.executor_pool = ProcessPoolExecutor(max_workers=3)
         self.executor_pool = ProcessPoolExecutor(max_workers=3)
         logger.info(self.executor_pool)
         logger.info(self.executor_pool)
@@ -105,6 +110,13 @@ class FrameTakeWorker:
         )
         )
 
 
         self.frame_gen = self.video_source.__iter__()
         self.frame_gen = self.video_source.__iter__()
+        self.frame_timeout = frame_timeout
+
+        self.pool_url = f"{pool_url}/api/pool"
+        self.auth = auth
+        self.push_timeout = push_timeout
+
+        self.bytes_queue = asyncio.Queue(maxsize=16)
 
 
     async def get_frame(self):
     async def get_frame(self):
         return next(self.frame_gen)
         return next(self.frame_gen)
@@ -115,24 +127,69 @@ class FrameTakeWorker:
 
 
     async def frame_extractor(self, frame_bytes: bytes):
     async def frame_extractor(self, frame_bytes: bytes):
         loop = asyncio.get_event_loop()
         loop = asyncio.get_event_loop()
-        logger.info(f"Input bytes: {len(frame_bytes)}")
+        logger.info(f"Take frame bytes: {len(frame_bytes)}")
         output = await loop.run_in_executor(self.executor_pool, video2_extractor, frame_bytes)
         output = await loop.run_in_executor(self.executor_pool, video2_extractor, frame_bytes)
         return output
         return output
 
 
     async def generator_worker(self):
     async def generator_worker(self):
+        await asyncio.sleep(random.uniform(1, 3))
         while True:
         while True:
-            logger.info("get frame")
+            get_frame = time.monotonic()
             frame = await self.get_frame()
             frame = await self.get_frame()
-            logger.info(f"{frame}")
             data = await self.frame_extractor(bytes(frame))
             data = await self.frame_extractor(bytes(frame))
-            logger.info(f"{len(data)}")
+            await self.bytes_queue.put(data)
+            put_frame = time.monotonic()
+            frame_time = get_frame - put_frame
+
+            if frame_time < self.frame_timeout:
+                await asyncio.sleep(math.ceil(self.frame_timeout - frame_time))
+            else:
+                await asyncio.sleep(random.uniform(1, 3))
+
+    async def push_worker(self):
+        await asyncio.sleep(self.push_timeout)
+        while True:
+            data_bytes = b""
+
+            for _ in range(self.bytes_queue.qsize()):
+                data_bytes += await self.bytes_queue.get()
+
+            async with aiohttp.ClientSession() as session:
+                async with session.post(self.pool_url, data=data_bytes, headers={"X-Secret": self.auth}) as response:
+                    if response.status == 200:
+                        logger.success(f"{response.status}: {response.text}")
+                    else:
+                        logger.error(f"{response.status}: {response.text}")
+            await asyncio.sleep(self.push_timeout)
 
 
     async def run_workers(self):
     async def run_workers(self):
         tasks = [self.generator_worker() for _ in range(self.processes)]
         tasks = [self.generator_worker() for _ in range(self.processes)]
+        tasks.append(self.push_worker())
         await asyncio.gather(*tasks)
         await asyncio.gather(*tasks)
-        pass
 
 
 
 
 if __name__ == "__main__":
 if __name__ == "__main__":
-    ftw = FrameTakeWorker(0, 3)
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source", type=str, required=True)
+    # parser.add_argument("--source-type", type=str, default="video")
+    parser.add_argument("--secret-file", type=str, default="./.secret")
+    parser.add_argument("--frame-timeout", type=int, default=1)
+    parser.add_argument("--pool-url", type=str, default="https://yebi.su")
+    parser.add_argument("--push-timeout", type=int, default=10)
+    parser.add_argument("--tmp-dir", type=str)
+
+    args = parser.parse_args()
+
+    with open(args.secret_file, "r") as f:
+        lines = f.read().strip().split("\n")
+        ident = lines[0].strip()
+        secret = lines[1].strip()
+
+        secret = f"{ident} {secret}"
+
+    ftw = GeneratorWorker(pool_url=args.pool_url, auth=secret,
+                          device_id=args.source,
+                          frame_timeout=args.frame_timeout,
+                          push_timeout=args.push_timeout)
+
     asyncio.run(ftw.run_workers())
     asyncio.run(ftw.run_workers())