nedo-vision-worker 1.0.0__py3-none-any.whl

nedo_vision_worker/services/SystemUsageClient.py

@@ -0,0 +1,77 @@
+from .GrpcClientBase import GrpcClientBase
+from ..protos.VisionWorkerService_pb2_grpc import VisionWorkerServiceStub
+from ..protos.VisionWorkerService_pb2 import SystemUsageRequest, GPUUsage
+import logging
+
+
+class SystemUsageClient(GrpcClientBase):
+    def __init__(self, server_host: str, server_port: int = 50051):
+        """
+        Initialize the system usage client.
+
+        Args:
+            server_host (str): The server hostname or IP address.
+            server_port (int): The server port. Default is 50051.
+        """
+        super().__init__(server_host, server_port)
+
+        try:
+            self.connect(VisionWorkerServiceStub)
+        except Exception as e:
+            logging.error(f"Failed to connect to gRPC server: {e}")
+            self.stub = None
+
+    def send_system_usage(self, device_id: str, cpu_usage: float, ram_usage: dict, gpu_usage: list, latency: float, token: str) -> dict:
+        """
+        Send system usage data to the server using token authentication.
+
+        Args:
+            device_id (str): The unique device ID.
+            cpu_usage (float): CPU usage percentage.
+            ram_usage (dict): RAM usage details.
+            gpu_usage (list): GPU usage details.
+            latency (float): Measured network latency in milliseconds.
+            token (str): Authentication token for the worker.
+
+        Returns:
+            dict: A dictionary containing the result of sending system usage.
+        """
+        if not self.stub:
+            return {"success": False, "message": "gRPC connection is not established."}
+
+        try:
+            # Prepare the request
+            request = SystemUsageRequest(
+                device_id=device_id,
+                cpu_usage=cpu_usage,
+                ram_usage_percent=ram_usage.get("percent", 0.0),
+                ram_total=ram_usage.get("total", 0),
+                ram_used=ram_usage.get("used", 0),
+                ram_free=ram_usage.get("free", 0),
+                latency_ms=latency, 
+                token=token,
+                gpu=[
+                    GPUUsage(
+                        gpu_index=gpu.get("gpu_index", 0),
+                        gpu_usage_percent=gpu.get("gpu_usage_percent", 0.0),
+                        memory_usage_percent=gpu.get("memory_usage_percent", 0.0),
+                        temperature_celsius=gpu.get("temperature_celsius", 0.0),
+                        total_memory=gpu.get("total_memory", 0),
+                        used_memory=gpu.get("used_memory", 0),
+                        free_memory=gpu.get("free_memory", 0),
+                    )
+                    for gpu in (gpu_usage or [])  # Handle None or empty list
+                ],
+            )
+
+            # Call the SendSystemUsage RPC
+            response = self.handle_rpc(self.stub.SendSystemUsage, request)
+
+            if response and response.success:
+                return {"success": True, "message": response.message}
+
+            return {"success": False, "message": response.message if response else "Unknown error"}
+
+        except Exception as e:
+            logging.error(f"Error sending system usage: {e}")
+            return {"success": False, "message": f"Error occurred: {e}"}

nedo_vision_worker/services/VideoStreamClient.py

@@ -0,0 +1,161 @@
+import grpc
+import time
+import logging
+import json
+import subprocess
+import ffmpeg
+import fractions
+from urllib.parse import urlparse
+from .GrpcClientBase import GrpcClientBase
+from ..protos.VisionWorkerService_pb2_grpc import VideoStreamServiceStub
+from ..protos.VisionWorkerService_pb2 import VideoFrame
+
+
+class VideoStreamClient(GrpcClientBase):
+    def __init__(self, server_host: str, server_port: int = 50051):
+        """Initialize the video stream client."""
+        super().__init__(server_host, server_port)
+
+    def _detect_stream_type(self, url):
+        """Detect whether the stream is RTSP or HLS based on the URL scheme."""
+        parsed_url = urlparse(url)
+        if parsed_url.scheme == "rtsp":
+            return "rtsp"
+        elif parsed_url.scheme in ["http", "https"] and url.endswith(".m3u8"):
+            return "hls"
+        else:
+            return "unknown"
+
+    def _get_video_properties(self, url, stream_type):
+        """Extract FPS, resolution, and pixel format dynamically for RTSP or HLS."""
+        try:
+            probe_cmd = [
+                "ffprobe",
+                "-i", url,
+                "-select_streams", "v:0",
+                "-show_entries", "stream=width,height,r_frame_rate,pix_fmt",
+                "-of", "json",
+                "-v", "quiet"
+            ]
+
+            if stream_type == "rtsp":
+                probe_cmd.insert(1, "-rtsp_transport")
+                probe_cmd.insert(2, "tcp")
+
+            result = subprocess.run(probe_cmd, capture_output=True, text=True)
+            probe_data = json.loads(result.stdout)
+
+            if "streams" in probe_data and len(probe_data["streams"]) > 0:
+                stream = probe_data["streams"][0]
+                width = int(stream["width"])
+                height = int(stream["height"])
+                fps = float(fractions.Fraction(stream["r_frame_rate"]))  # Safe FPS conversion
+                pixel_format = stream["pix_fmt"]
+                return width, height, fps, pixel_format
+
+        except Exception as e:
+            logging.error(f"Error extracting video properties: {e}")
+
+        return None, None, None, None
+
+    def _get_bytes_per_pixel(self, pixel_format):
+        """Determine bytes per pixel based on pixel format."""
+        pixel_map = {"rgb24": 3, "yuv420p": 1.5, "gray": 1}
+        return pixel_map.get(pixel_format, 3)  # Default to 3 bytes (RGB)
+
+    def _generate_frames(self, url, worker_id, uuid, stream_duration):
+        """Generator function to continuously stream video frames to gRPC."""
+        stream_type = self._detect_stream_type(url)
+        if stream_type == "unknown":
+            logging.error(f"Unsupported stream type: {url}")
+            return
+
+        width, height, fps, pixel_format = self._get_video_properties(url, stream_type)
+        if not width or not height or not fps:
+            logging.error("Failed to retrieve stream properties.")
+            return
+
+        bytes_per_pixel = self._get_bytes_per_pixel(pixel_format)
+        frame_size = int(width * height * bytes_per_pixel)
+        frame_interval = 1.0 / fps  # Time between frames
+        start_time = time.time()
+        empty_frame_count = 0
+
+        logging.info(f"Streaming {stream_type.upper()} from: {url} for {stream_duration} seconds...")
+
+        if stream_type == "rtsp":
+            ffmpeg_input = (
+                ffmpeg
+                .input(url, rtsp_transport="tcp", fflags="nobuffer", timeout="5000000")
+            )
+        elif stream_type == "hls":
+            ffmpeg_input = (
+                ffmpeg
+                .input(url, format="hls", analyzeduration="10000000", probesize="10000000")
+            )
+        else:
+            logging.error(f"Unsupported stream type: {url}")
+            return
+
+        process = (
+            ffmpeg_input
+            .output("pipe:", format="rawvideo", pix_fmt=pixel_format, vsync="passthrough")
+            .overwrite_output()  # Replaces `global_args()` for avoiding conflicts
+            .run_async(pipe_stdout=True, pipe_stderr=True)
+        )
+
+        try:
+            while time.time() - start_time < stream_duration:
+                frame_bytes = process.stdout.read(frame_size)
+
+                if not frame_bytes:
+                    empty_frame_count += 1
+                    logging.warning(f"Empty frame received ({empty_frame_count}), retrying...")
+
+                    if empty_frame_count > 5:  # Stop if 5 consecutive empty frames
+                        logging.error("Too many empty frames, stopping stream...")
+                        break
+                    continue  # Try reading the next frame
+
+                empty_frame_count = 0  # Reset empty frame count
+                yield VideoFrame(
+                    worker_id=worker_id,
+                    uuid=uuid,
+                    frame_data=frame_bytes,
+                    timestamp=int(time.time() * 1000),
+                )
+
+                time.sleep(frame_interval)  # Ensure proper frame timing
+
+        except Exception as e:
+            logging.error(f"Streaming error: {e}")
+
+        finally:
+            stderr_output = process.stderr.read().decode()
+            logging.error(f"FFmpeg stderr: {stderr_output}")  # Log errors from FFmpeg
+            process.terminate()
+            process.wait()  # Ensures cleanup
+
+    def stream_video(self, worker_id, uuid, url, stream_duration):
+        """
+        Stream video frames from RTSP or HLS to gRPC server.
+
+        Args:
+            worker_id (str): Worker ID
+            uuid (str): Unique stream session ID
+            url (str): Stream URL (RTSP or HLS)
+            stream_duration (int): Duration in seconds to stream
+        """
+        self.connect(VideoStreamServiceStub)  # Ensure connection and stub are established
+
+        try:
+            for response in self.stub.StreamVideo(self._generate_frames(url, worker_id, uuid, stream_duration)):
+                if response.success:
+                    logging.info(f"Frame sent successfully: {response.message}")
+                else:
+                    logging.error(f"Frame rejected: {response.message}")
+
+        except grpc.RpcError as e:
+            logging.error(f"gRPC error: {e.code()} - {e.details()}")  # Log more details
+        except Exception as e:
+            logging.error(f"Unexpected streaming error: {e}")

nedo_vision_worker/services/WorkerSourceClient.py

@@ -0,0 +1,215 @@
+import logging
+import os
+from pathlib import Path
+from ..database.DatabaseManager import _get_storage_paths
+from ..models.worker_source import WorkerSourceEntity
+from .GrpcClientBase import GrpcClientBase
+from ..protos.WorkerSourceService_pb2_grpc import WorkerSourceServiceStub
+from ..protos.WorkerSourceService_pb2 import (
+    GetWorkerSourceListRequest,
+    UpdateWorkerSourceRequest,
+    DownloadSourceFileRequest
+)
+from ..repositories.WorkerSourceRepository import WorkerSourceRepository
+
+logger = logging.getLogger(__name__)
+
+class WorkerSourceClient(GrpcClientBase):
+    def __init__(self, server_host: str, server_port: int = 50051):
+        super().__init__(server_host, server_port)
+        storage_paths = _get_storage_paths()
+        self.source_file_path = storage_paths["files"] / "source_files"
+        self.source_file_path.mkdir(parents=True, exist_ok=True)
+        self.repo = WorkerSourceRepository()
+        try:
+            self.connect(WorkerSourceServiceStub)
+        except Exception as e:
+            logging.error(f"Failed to connect to gRPC server: {e}")
+            self.stub = None
+
+    def _get_path(self, file: str) -> Path:
+        """Get the path to a local AI model file."""
+        return self.source_file_path / os.path.basename(file)
+
+    def sync_worker_sources(self, worker_id: str, token: str) -> dict:
+        """Fetch and sync worker source list from gRPC service using token authentication."""
+        if not self.stub:
+            return {"success": False, "message": "gRPC connection is not established."}
+
+        try:
+            request = GetWorkerSourceListRequest(worker_id=worker_id, token=token)
+            response = self.handle_rpc(self.stub.GetWorkerSourceList, request)
+
+            if response and response.success:
+                self._process_server_sources(response.data)
+                return {"success": True, "message": response.message, "data": response.data}
+
+            return {"success": False, "message": response.message if response else "Unknown error"}
+
+        except Exception as e:
+            logging.error(f"Error fetching worker source list: {e}")
+            return {"success": False, "message": f"Error occurred: {e}"}
+        
+    def _process_server_sources(self, server_sources):
+        """Process server sources, handling additions, updates, and deletions."""
+        local_sources = {source.id: source for source in self.repo.get_all_worker_sources()}
+        server_source_ids = set()
+        
+        new_sources = []
+        updated_sources = []
+        changed_records = []
+        
+        # Process each source from the server
+        for source in server_sources:
+            server_source_ids.add(source.id)
+            existing_source = local_sources.get(source.id)
+
+            if existing_source:
+                self._handle_existing_source(source, existing_source, updated_sources, changed_records)
+            else:
+                self._handle_new_source(source, new_sources)
+        
+        # Handle sources that no longer exist on the server
+        sources_to_delete = [
+            source for source_id, source in local_sources.items()
+            if source_id not in server_source_ids
+        ]
+        
+        self._save_changes(new_sources, updated_sources, sources_to_delete, changed_records)
+
+    def _handle_existing_source(self, source, existing_source, updated_sources, changed_records):
+        """Handle source that exists locally but might need updates."""
+        changes = []
+
+        # Check for field changes and log them
+        if existing_source.name != source.name:
+            changes.append(f"name: '{existing_source.name}' → '{source.name}'")
+        if existing_source.worker_id != source.worker_id:
+            changes.append(f"worker_id: {existing_source.worker_id} → {source.worker_id}")
+        if existing_source.type_code != source.type_code:
+            if source.type_code == "live":
+                self.delete_local_source_file(existing_source.file_path)
+            elif existing_source.file_path == source.file_path:
+                self.download_source_file(source)
+
+            changes.append(f"type_code: {existing_source.type_code} → {source.type_code}")
+        if existing_source.url != source.url:
+            changes.append(f"url: {existing_source.url} → {source.url}")
+        if existing_source.file_path != source.file_path:
+            self.delete_local_source_file(existing_source.file_path)
+            self.download_source_file(source)
+
+            changes.append(f"file_path: {existing_source.file_path} → {source.file_path}")
+
+        if changes:
+            # Update existing record
+            existing_source.name = source.name
+            existing_source.worker_id = source.worker_id
+            existing_source.type_code = source.type_code
+            existing_source.url = source.url
+            existing_source.file_path = source.file_path
+
+            updated_sources.append(existing_source)
+            changed_records.append(f"🔄 [APP] [UPDATE] Worker Source ID {source.id}: " + ", ".join(changes))
+
+    def _handle_new_source(self, source, new_sources):
+        """Handle source that doesn't exist locally."""
+        new_record = WorkerSourceEntity(
+                        id=source.id,
+                        name=source.name,
+                        worker_id=source.worker_id,
+                        type_code=source.type_code,
+                        file_path=source.file_path,
+                        url=source.url
+                    )
+        new_sources.append(new_record)
+        logger.info(f"🆕 [APP] [INSERT] Added Worker Source ID {source.id} - {source.name}")
+        if source.type_code == "file":
+            self.download_source_file(source)
+
+    def _save_changes(self, new_sources, updated_sources, sources_to_delete, changed_records):
+        """Save all changes to database in a single transaction."""
+        try:
+            if new_sources:
+                self.repo.session.bulk_save_objects(new_sources)
+
+            if updated_sources:
+                self.repo.session.bulk_save_objects(updated_sources)
+                
+            for source in sources_to_delete:
+                logger.info(f"❌ [APP] [DELETE] Removing Worker Source ID {source.id} - {source.name}")
+                self.repo.session.delete(source)
+                self.delete_local_source_file(source.file_path)
+                
+            self.repo.session.commit()
+
+            for change in changed_records:
+                logger.info(change)
+
+        except Exception as e:
+            self.repo.session.rollback()
+            logging.error(f"Error saving sources changes: {e}")
+            raise
+
+    def update_worker_source(
+        self, worker_source_id: str, resolution: str, status_code: str, frame_rate: float, worker_timestamp: str, token: str
+    ) -> dict:
+        """Updates a worker source entry using gRPC with token authentication."""
+        if not self.stub:
+            return {"success": False, "message": "gRPC connection is not established."}
+
+        try:
+            request = UpdateWorkerSourceRequest(
+                worker_source_id=worker_source_id,
+                resolution=resolution,
+                status_code=status_code,
+                frame_rate=frame_rate,
+                worker_timestamp=worker_timestamp,
+                token=token
+            )
+            response = self.handle_rpc(self.stub.Update, request)
+
+            if response and response.success:
+                return {"success": True, "message": response.message}
+
+            return {"success": False, "message": response.message if response else "Unknown error"}
+
+        except Exception as e:
+            logging.error(f"Error updating worker source: {e}")
+            return {"success": False, "message": f"Error occurred: {e}"}
+
+    def download_source_file(self, source) -> bool:
+        """Download the AI model and save it to the models directory."""
+        if not self.stub:
+            logging.error("gRPC connection is not established.")
+            return False
+        if not source.file_path:
+            return False
+        
+        try:
+            logging.info(f"📥 Downloading source file '{source.name}'...")
+            request = DownloadSourceFileRequest(source_id=source.id)
+            file_path = self._get_path(source.file_path)
+            
+            with open(file_path, "wb") as f:
+                for chunk in self.stub.DownloadSourceFile(request):
+                    f.write(chunk.file_chunk)
+            
+            logging.info(f"✅ Source File '{source.name}' downloaded successfully")
+            return True
+        
+        except Exception as e:
+            logging.error(f"❌ Error downloading Source File '{source.name}': {e}")
+            return False
+        
+    def delete_local_source_file(self, file: str) -> None:
+        """Delete a local Source file."""
+        if not file:
+            return
+
+        file_path = self._get_path(file)
+        try:
+            file_path.unlink()
+            logging.info(f"🗑️ File deleted")
+        except Exception as e:
+            logging.error(f"❌ Error deleting file: {e}")