nedo-vision-worker-core 0.2.0__py3-none-any.whl

nedo_vision_worker_core/models/__init__.py

@@ -0,0 +1,20 @@
+# Import all models to ensure they are registered with SQLAlchemy Base registry
+# This prevents the auto-migration system from thinking these models don't exist
+
+from .ai_model import AIModelEntity
+from .auth import AuthEntity
+from .config import ConfigEntity
+from .dataset_source import DatasetSourceEntity
+from .logs import LogEntity
+from .ppe_detection import PPEDetectionEntity
+from .ppe_detection_label import PPEDetectionLabelEntity
+from .restricted_area_violation import RestrictedAreaViolationEntity
+from .user import UserEntity
+from .worker_source import WorkerSourceEntity
+from .worker_source_pipeline import WorkerSourcePipelineEntity
+from .worker_source_pipeline_config import WorkerSourcePipelineConfigEntity
+from .worker_source_pipeline_debug import WorkerSourcePipelineDebugEntity
+from .worker_source_pipeline_detection import WorkerSourcePipelineDetectionEntity
+
+# This ensures all models are imported and registered with SQLAlchemy's Base registry
+# when the models package is imported during database initialization 

nedo_vision_worker_core/models/ai_model.py

@@ -0,0 +1,41 @@
+import uuid
+from sqlalchemy import Column, String, DateTime
+from datetime import datetime
+from ..database.DatabaseManager import Base
+
+class AIModelEntity(Base):
+    __tablename__ = "ai_model"
+    __bind_key__ = "default"
+
+    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+    file = Column(String, nullable=False)
+    type = Column(String, nullable=False)
+    name = Column(String, nullable=False)
+    version = Column(String, nullable=False)
+    download_status = Column(String, nullable=True, default="completed")  # pending, downloading, completed, failed
+    last_download_attempt = Column(DateTime, nullable=True)
+    download_error = Column(String, nullable=True)
+
+    def __repr__(self):
+        return (
+            f"<AIModelEntity(id={self.id}, name={self.name}, type={self.type}, "
+            f"file={self.file}, version={self.version})>"
+        )
+
+    def __str__(self):
+        return (
+            f"AIModelEntity(id={self.id}, name={self.name}, type={self.type}, "
+            f"file={self.file}, version={self.version}, status={self.download_status})"
+        )
+
+    def is_ready_for_use(self) -> bool:
+        """Check if the model is ready for use (downloaded and available)."""
+        return self.download_status == "completed"
+
+    def is_downloading(self) -> bool:
+        """Check if the model is currently being downloaded."""
+        return self.download_status in ["pending", "downloading"]
+
+    def has_download_failed(self) -> bool:
+        """Check if the model download has failed."""
+        return self.download_status == "failed"

nedo_vision_worker_core/models/auth.py

@@ -0,0 +1,14 @@
+from sqlalchemy import Column, String, Integer
+from ..database.DatabaseManager import Base
+
+class AuthEntity(Base):
+    __tablename__ = "auth"
+    id = Column(Integer, primary_key=True)
+    username = Column(String, nullable=False)
+    password = Column(String, nullable=False)
+
+    def __repr__(self):
+        return f"<AuthEntity(id={self.id}, username={self.username})>"
+
+    def to_dict(self):
+        return {"id": self.id, "username": self.username}

nedo_vision_worker_core/models/config.py

@@ -0,0 +1,9 @@
+from sqlalchemy import Column, String
+from ..database.DatabaseManager import Base
+
+class ConfigEntity(Base):
+    __tablename__ = "server_config"
+    __bind_key__ = "config"
+
+    key = Column(String, primary_key=True)
+    value = Column(String, nullable=False)

nedo_vision_worker_core/models/dataset_source.py

@@ -0,0 +1,30 @@
+from sqlalchemy import Column, String, Integer
+from ..database.DatabaseManager import Base
+
+class DatasetSourceEntity(Base):
+    __tablename__ = "dataset_sources"
+    __bind_key__ = "default"
+
+    id = Column(String, primary_key=True)
+    dataset_id = Column(String, nullable=False)
+    worker_source_id = Column(String, nullable=False)
+    sampling_interval = Column(Integer, nullable=False)
+    dataset_name = Column(String, nullable=False)
+    worker_source_name = Column(String, nullable=False)
+    worker_source_url = Column(String, nullable=False)
+
+    def __repr__(self):
+        return (
+            f"<DatasetSourceEntity(id={self.id}, dataset_id={self.dataset_id}, "
+            f"worker_source_id={self.worker_source_id}, sampling_interval={self.sampling_interval}, "
+            f"dataset_name={self.dataset_name}, worker_source_name={self.worker_source_name}, "
+            f"worker_source_url={self.worker_source_url})>"
+        )
+
+    def __str__(self):
+        return (
+            f"DatasetSourceEntity(id={self.id}, dataset_id={self.dataset_id}, "
+            f"worker_source_id={self.worker_source_id}, sampling_interval={self.sampling_interval}, "
+            f"dataset_name={self.dataset_name}, worker_source_name={self.worker_source_name}, "
+            f"worker_source_url={self.worker_source_url})"
+        ) 

nedo_vision_worker_core/models/logs.py

@@ -0,0 +1,9 @@
+from sqlalchemy import Column, String, Integer
+from ..database.DatabaseManager import Base
+
+class LogEntity(Base):
+    __tablename__ = "logs"
+    __bind_key__ = "logging"
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    message = Column(String, nullable=False)

nedo_vision_worker_core/models/ppe_detection.py

@@ -0,0 +1,39 @@
+import uuid
+import datetime
+from sqlalchemy import Column, String, ForeignKey, DateTime, Float, Integer
+from sqlalchemy.orm import relationship
+from ..database.DatabaseManager import Base
+
+class PPEDetectionEntity(Base):
+    __tablename__ = "ppe_detections"
+    __bind_key__ = "default"
+
+    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+    worker_id = Column(String, nullable=False)
+    worker_source_id = Column(String, nullable=False)
+    person_id = Column(String, nullable=False)
+    image_path = Column(String, nullable=False)
+    image_tile_path = Column(String, nullable=False)
+    detection_count = Column(Integer, nullable=False, default=0)  # Tracks total detections before saving
+    created_at = Column(DateTime, default=datetime.datetime.utcnow)
+    b_box_x1 = Column(Float, nullable=False)
+    b_box_y1 = Column(Float, nullable=False)
+    b_box_x2 = Column(Float, nullable=False)
+    b_box_y2 = Column(Float, nullable=False)
+
+    ppe_labels = relationship("PPEDetectionLabelEntity", back_populates="detection")
+
+    def __repr__(self):
+        return (
+            f"<PPEDetectionEntity(id={self.id}, worker_id={self.worker_id}, "
+            f"worker_source_id={self.worker_source_id}, person_id={self.person_id}, "
+            f"image_path={self.image_path}, detection_count={self.detection_count}, "
+            f"created_at={self.created_at})>"
+        )
+
+    def __str__(self):
+        return (
+            f"PPEDetectionEntity(id={self.id}, worker_id={self.worker_id}, "
+            f"worker_source_id={self.worker_source_id}, person_id={self.person_id}, "
+            f"detection_count={self.detection_count}, created_at={self.created_at})"
+        )

nedo_vision_worker_core/models/ppe_detection_label.py

@@ -0,0 +1,20 @@
+import uuid
+from sqlalchemy import Column, String, ForeignKey, DateTime, Float, Integer
+from sqlalchemy.orm import relationship
+from ..database.DatabaseManager import Base
+
+class PPEDetectionLabelEntity(Base):
+    __tablename__ = "ppe_detection_labels"
+    __bind_key__ = "default"
+
+    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+    detection_id = Column(String, ForeignKey("ppe_detections.id"), nullable=False)
+    code = Column(String, nullable=False)  # helmet, vest, etc.
+    confidence_score = Column(Float, nullable=False)
+    detection_count = Column(Integer, nullable=False, default=0)
+    b_box_x1 = Column(Float, nullable=False)
+    b_box_y1 = Column(Float, nullable=False)
+    b_box_x2 = Column(Float, nullable=False)
+    b_box_y2 = Column(Float, nullable=False)
+
+    detection = relationship("PPEDetectionEntity", back_populates="ppe_labels")

nedo_vision_worker_core/models/restricted_area_violation.py

@@ -0,0 +1,20 @@
+import uuid
+import datetime
+from sqlalchemy import Column, Float, String, DateTime
+from ..database.DatabaseManager import Base
+
+class RestrictedAreaViolationEntity(Base):
+    __tablename__ = "restricted_area_violation"
+    __bind_key__ = "default"
+
+    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+    worker_source_id = Column(String, nullable=False)
+    person_id = Column(String, nullable=False)
+    image_path = Column(String, nullable=False)
+    image_tile_path = Column(String, nullable=False)
+    confidence_score = Column(Float, nullable=False)
+    created_at = Column(DateTime, default=datetime.datetime.utcnow)
+    b_box_x1 = Column(Float, nullable=False)
+    b_box_y1 = Column(Float, nullable=False)
+    b_box_x2 = Column(Float, nullable=False)
+    b_box_y2 = Column(Float, nullable=False)

nedo_vision_worker_core/models/user.py

@@ -0,0 +1,10 @@
+from sqlalchemy import Column, String
+from ..database.DatabaseManager import Base
+
+class UserEntity(Base):
+    __tablename__ = "user"
+    __bind_key__ = "auth"
+
+    id = Column(String, primary_key=True)
+    username = Column(String, nullable=False)
+    password = Column(String, nullable=False)

nedo_vision_worker_core/models/worker_source.py

@@ -0,0 +1,19 @@
+from sqlalchemy import Column, String, Float
+from ..database.DatabaseManager import Base
+
+class WorkerSourceEntity(Base):
+    __tablename__ = "worker_source"
+    __bind_key__ = "config"
+
+    id = Column(String, primary_key=True)
+    name = Column(String, nullable=False)
+    worker_id = Column(String, nullable=False)
+    type_code = Column(String, nullable=False)
+    file_path = Column(String, nullable=True)
+    url = Column(String, nullable=False)
+    resolution = Column(String, nullable=True)
+    status_code = Column(String, nullable=True)
+    frame_rate = Column(Float, nullable=True)
+    source_location_code = Column(String, nullable=True)  # Optional field
+    latitude = Column(Float, nullable=True)  # Optional field
+    longitude = Column(Float, nullable=True)  # Optional field

nedo_vision_worker_core/models/worker_source_pipeline.py

@@ -0,0 +1,21 @@
+from sqlalchemy import Column, String
+from sqlalchemy.orm import relationship
+from ..database.DatabaseManager import Base
+
+class WorkerSourcePipelineEntity(Base):
+    __tablename__ = "worker_source_pipeline"
+    __bind_key__ = "config"
+
+    id = Column(String, primary_key=True)
+    name = Column(String, nullable=False)
+    worker_source_id = Column(String, nullable=False)
+    worker_id = Column(String, nullable=False)
+    ai_model_id = Column(String, nullable=True)
+    pipeline_status_code = Column(String, nullable=False)
+    location_name = Column(String, nullable=True)
+
+    worker_source_pipeline_configs = relationship(
+        "WorkerSourcePipelineConfigEntity",
+        back_populates="pipeline",
+        cascade="all, delete-orphan"
+    )

nedo_vision_worker_core/models/worker_source_pipeline_config.py

@@ -0,0 +1,24 @@
+from sqlalchemy import Column, String, ForeignKey, Boolean
+from sqlalchemy.orm import relationship
+from ..database.DatabaseManager import Base
+
+
+class WorkerSourcePipelineConfigEntity(Base):
+    __tablename__ = "worker_source_pipeline_config"
+    __bind_key__ = "config"
+
+    id = Column(String, primary_key=True)
+    worker_source_pipeline_id = Column(
+        String, ForeignKey("worker_source_pipeline.id", ondelete="CASCADE"), nullable=False
+    )
+    pipeline_config_id = Column(String, nullable=False)
+    is_enabled = Column(Boolean, nullable=False)
+    value = Column(String, nullable=True)
+    pipeline_config_name = Column(String, nullable=False)
+    pipeline_config_code = Column(String, nullable=False)
+
+    pipeline = relationship(
+        "WorkerSourcePipelineEntity",
+        back_populates="worker_source_pipeline_configs",
+        passive_deletes=True 
+    )

nedo_vision_worker_core/models/worker_source_pipeline_debug.py

@@ -0,0 +1,15 @@
+from datetime import datetime
+import uuid
+from sqlalchemy import Column, DateTime, String
+from ..database.DatabaseManager import Base
+
+class WorkerSourcePipelineDebugEntity(Base):
+    __tablename__ = "worker_source_pipeline_debug"
+    __bind_key__ = "default"
+    
+    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+    uuid = Column(String, nullable=False)
+    worker_source_pipeline_id = Column(String, nullable=False)
+    image_path = Column(String, nullable=True)
+    data = Column(String, nullable=True)
+    created_at = Column(DateTime, default=datetime.utcnow)

nedo_vision_worker_core/models/worker_source_pipeline_detection.py

@@ -0,0 +1,14 @@
+from datetime import datetime
+import uuid
+from sqlalchemy import Column, DateTime, String
+from ..database.DatabaseManager import Base
+
+class WorkerSourcePipelineDetectionEntity(Base):
+    __tablename__ = "worker_source_pipeline_detection"
+    __bind_key__ = "default"
+    
+    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+    worker_source_pipeline_id = Column(String, nullable=False)
+    image_path = Column(String, nullable=True)
+    data = Column(String, nullable=True)
+    created_at = Column(DateTime, default=datetime.utcnow)

nedo_vision_worker_core/pipeline/PipelineConfigManager.py

@@ -0,0 +1,32 @@
+from ..repositories.WorkerSourcePipelineRepository import WorkerSourcePipelineRepository
+
+
+class PipelineConfigManager:
+    def __init__(self):
+        self.config_repository = WorkerSourcePipelineRepository()
+        self.config = {}
+    
+    def update(self, pipeline_id):
+        self.config = self.config_repository.get_pipeline_configs_by_pipeline_id(pipeline_id)
+
+    def is_feature_enabled(self, feature_name):
+        """
+        Generic method to check if a feature is enabled in configuration.
+        
+        :param feature_name: Name of the feature to check
+        :return: True if feature is enabled, False otherwise
+        """
+        return (feature_name in self.config and 
+                self.config[feature_name].get("is_enabled", False))
+    
+    def get_feature_config(self, feature_name, default = {}) -> dict:
+        """
+        Get the configuration for a specific feature.
+        
+        :param feature_name: Name of the feature
+        :return : Feature configuration
+        """
+        if not self.is_feature_enabled(feature_name):
+            return default
+
+        return self.config[feature_name].get("value", default)

nedo_vision_worker_core/pipeline/PipelineManager.py

@@ -0,0 +1,133 @@
+import logging
+import threading
+from concurrent.futures import ThreadPoolExecutor, Future
+from typing import Dict
+from .PipelineProcessor import PipelineProcessor
+from ..streams.VideoStreamManager import VideoStreamManager
+
+class PipelineManager:
+    """Manages AI pipeline execution and video stream processing."""
+
+    def __init__(self, video_manager: VideoStreamManager, on_pipeline_stopped, max_workers=50):
+        self.executor = ThreadPoolExecutor(max_workers=max_workers)  # Thread pool for parallel execution
+        self.pipeline_threads = {}  # Stores Future objects {pipeline_id: Future}
+        self.pipeline_metadata = {}  # Stores actual pipeline data {pipeline_id: metadata}
+        self.video_manager = video_manager  # Manages video streams
+        self.processors: Dict[str, PipelineProcessor] = {}  # Stores PipelineProcessor instances per pipeline
+        self.running = True
+        self._stopping_pipelines = set()  # Track pipelines being stopped
+        self._stop_lock = threading.Lock()  # Lock for thread-safe pipeline stopping
+        self.on_pipeline_stopped = on_pipeline_stopped
+
+    def start_pipeline(self, pipeline, model):
+        """
+        Start a pipeline processing.
+        Args:
+            pipeline: The pipeline object (contains id, worker_source_id, name, etc.)
+            model: The AI model to use for processing.
+        """
+        pipeline_id = pipeline.id
+        worker_source_id = pipeline.worker_source_id
+
+        if not self.running:
+            logging.warning(f"⚠️ Attempt to start pipeline {pipeline_id} after shutdown.")
+            return
+
+        if self.is_running(pipeline_id):
+            logging.warning(f"⚠️ Pipeline {pipeline_id} is already running.")
+            return
+
+        logging.info(f"🚀 Starting Pipeline processing for pipeline: {pipeline_id} | Source: {worker_source_id} ({pipeline.name})")
+
+        processor = PipelineProcessor(pipeline_id, worker_source_id, model, False)
+        processor.frame_drawer.location_name = pipeline.location_name
+        self.processors[pipeline_id] = processor  # Store processor instance
+
+        future = self.executor.submit(processor.process_pipeline, self.video_manager)
+        self.pipeline_threads[pipeline_id] = future
+        self.pipeline_metadata[pipeline_id] = pipeline
+
+        # Add callback to detect when a pipeline finishes
+        future.add_done_callback(lambda f: self._handle_pipeline_completion(pipeline_id, f))
+
+    def _handle_pipeline_completion(self, pipeline_id: str, future: Future):
+        """
+        Handles cleanup when a pipeline finishes processing.
+        """
+        with self._stop_lock:
+            if pipeline_id in self._stopping_pipelines:
+                return  # If it's already being stopped manually, don't trigger again
+
+        try:
+            if future.cancelled():
+                logging.info(f"🚫 Pipeline {pipeline_id} was cancelled.")
+            elif future.exception():
+                logging.error(f"❌ Pipeline {pipeline_id} encountered an error: {future.exception()}", exc_info=True)
+
+        except Exception as e:
+            logging.error(f"⚠️ Error in handling pipeline {pipeline_id} completion: {e}")
+
+        finally:
+            self.on_pipeline_stopped(pipeline_id)
+
+    def stop_pipeline(self, pipeline_id: str):
+        """Stop an AI processing pipeline."""
+        with self._stop_lock:
+            if pipeline_id in self._stopping_pipelines:
+                logging.debug(f"Pipeline {pipeline_id} already being stopped, skipping")
+                return
+            self._stopping_pipelines.add(pipeline_id)
+
+        try:
+            # Stop AI processing
+            processor = self.processors.pop(pipeline_id, None)
+            if processor:
+                processor.stop()
+
+            # Cancel execution thread
+            future = self.pipeline_threads.pop(pipeline_id, None)
+            if future:
+                future.cancel()
+
+            # Remove metadata
+            self.pipeline_metadata.pop(pipeline_id, None)
+
+            logging.info(f"✅ Pipeline {pipeline_id} stopped successfully.")
+
+        except Exception as e:
+            logging.error(f"❌ Error during pipeline shutdown: {e}")
+        
+        finally:
+            self._stopping_pipelines.discard(pipeline_id)
+            self.on_pipeline_stopped(pipeline_id)
+
+    def get_active_pipelines(self):
+        """Returns a list of active pipeline IDs."""
+        return list(self.pipeline_metadata.keys())
+
+    def get_pipeline(self, pipeline_id):
+        """Returns the actual pipeline metadata (not the Future object)."""
+        return self.pipeline_metadata.get(pipeline_id, None)
+
+    def is_running(self, pipeline_id):
+        """
+        Checks if a pipeline is currently running.
+        
+        Args:
+            pipeline_id (str): The ID of the pipeline to check.
+            
+        Returns:
+            bool: True if the pipeline is running, False otherwise.
+        """
+        return pipeline_id in self.pipeline_threads and not self.pipeline_threads[pipeline_id].done()
+
+    def shutdown(self):
+        """Shuts down the pipeline manager gracefully."""
+        logging.info("🛑 Shutting down PipelineManager...")
+        self.running = False
+
+        for pipeline_id in list(self.pipeline_threads.keys()):
+            self.stop_pipeline(pipeline_id)
+
+        self.executor.shutdown(wait=True)  # Wait for all threads to finish
+        logging.info("✅ PipelineManager stopped.")

nedo_vision_worker_core/pipeline/PipelinePrepocessor.py

@@ -0,0 +1,40 @@
+from typing import List, Tuple
+import numpy as np
+from .PipelineConfigManager import PipelineConfigManager
+from ..preprocessing.ImageResizer import ImageResizer
+from ..preprocessing.ImageRoi import ImageRoi
+from ..preprocessing.Preprocessor import Preprocessor
+
+
+class PipelinePrepocessor:
+    def __init__(self):
+        self.preprocessors: List[Preprocessor] = [
+            ImageRoi(),
+            ImageResizer()
+        ]
+
+    def update(self, config: PipelineConfigManager):
+        for preprocessor in self.preprocessors:
+            preprocessor.update_config(config)
+    
+    def apply(self, image: np.ndarray) -> np.ndarray:
+        image = image.copy()
+
+        for preprocessor in self.preprocessors:
+            image = preprocessor.apply(image)
+
+        return image
+    
+    def revert_detections_bboxes(self, detections: list, dimension: Tuple[int, int]) -> np.ndarray:
+        if not detections or len(detections) < 1:
+            return detections
+        
+        bboxes = np.array([det["bbox"] for det in detections])
+
+        for preprocessor in reversed(self.preprocessors):
+            bboxes = preprocessor.revert_bboxes(bboxes, dimension)
+
+        for det, bbox in zip(detections, bboxes):
+            det["bbox"] = bbox
+
+        return detections