nedo-vision-worker-core 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

nedo_vision_worker_core/repositories/RestrictedAreaRepository.py

@@ -5,18 +5,17 @@ import cv2
 import datetime
 import logging
 from pathlib import Path
-from sqlalchemy.orm import Session
 from sqlalchemy.exc import SQLAlchemyError
+from .BaseRepository import BaseRepository
 from ..models.restricted_area_violation import RestrictedAreaViolationEntity
 from ..database.DatabaseManager import DatabaseManager
 from ..util.DrawingUtils import DrawingUtils
 
-class RestrictedAreaRepository:
+class RestrictedAreaRepository(BaseRepository):
     def __init__(self):
+        super().__init__(db_name="default")
         self.storage_dir = DatabaseManager.STORAGE_PATHS["files"] / "restricted_violations"
         os.makedirs(self.storage_dir, exist_ok=True)
-        self.db_manager = DatabaseManager()
-        self.session: Session = self.db_manager.get_session("default")
 
     def save_area_violation(self, pipeline_id, worker_source_id, frame_id, tracked_objects, frame, frame_drawer):
         """
@@ -50,45 +49,45 @@ class RestrictedAreaRepository:
             cv2.imwrite(cropped_image_path, cropped_image)
 
             try:
-                new_detection = RestrictedAreaViolationEntity(
-                    worker_source_id=worker_source_id,
-                    person_id=person_id,
-                    image_path=full_image_path,
-                    image_tile_path=cropped_image_path,
-                    confidence_score=tracked_obj.get("confidence", 1),
-                    b_box_x1=bbox[0],
-                    b_box_y1=bbox[1],
-                    b_box_x2=bbox[2],
-                    b_box_y2=bbox[3],
-                )
-                self.session.add(new_detection)
-                self.session.flush()
-                self.session.commit()
-                logging.info(f"✅ Inserted restricted area violation for Person {person_id}")
-
-                # Trigger detection callback
-                try:
-                    from ..core_service import CoreService
-                    from ..detection.detection_processing.HumanDetectionProcessor import HumanDetectionProcessor
-                    
-                    # Create unified detection data using the processor's factory method
-                    unified_data = HumanDetectionProcessor.create_detection_data(
-                        pipeline_id=pipeline_id,
+                with self._get_session() as session:
+                    new_detection = RestrictedAreaViolationEntity(
                         worker_source_id=worker_source_id,
                         person_id=person_id,
-                        detection_id=new_detection.id if hasattr(new_detection, 'id') else f"area_{person_id}_{current_datetime}",
-                        tracked_obj=tracked_obj,
                         image_path=full_image_path,
                         image_tile_path=cropped_image_path,
-                        frame_id=frame_id
+                        confidence_score=tracked_obj.get("confidence", 1),
+                        b_box_x1=bbox[0],
+                        b_box_y1=bbox[1],
+                        b_box_x2=bbox[2],
+                        b_box_y2=bbox[3],
                     )
-                    
-                    # Trigger callbacks
-                    CoreService.trigger_detection(unified_data)
-                    
-                except Exception as e:
-                    logging.warning(f"⚠️ Failed to trigger area violation callback: {e}")
+                    session.add(new_detection)
+                    session.flush()
+                    # Commit happens automatically via context manager
+                    logging.info(f"✅ Inserted restricted area violation for Person {person_id}")
+
+                    # Trigger detection callback
+                    try:
+                        from ..core_service import CoreService
+                        from ..detection.detection_processing.HumanDetectionProcessor import HumanDetectionProcessor
+                        
+                        # Create unified detection data using the processor's factory method
+                        unified_data = HumanDetectionProcessor.create_detection_data(
+                            pipeline_id=pipeline_id,
+                            worker_source_id=worker_source_id,
+                            person_id=person_id,
+                            detection_id=new_detection.id if hasattr(new_detection, 'id') else f"area_{person_id}_{current_datetime}",
+                            tracked_obj=tracked_obj,
+                            image_path=full_image_path,
+                            image_tile_path=cropped_image_path,
+                            frame_id=frame_id
+                        )
+                        
+                        # Trigger callbacks
+                        CoreService.trigger_detection(unified_data)
+                        
+                    except Exception as e:
+                        logging.warning(f"⚠️ Failed to trigger area violation callback: {e}")
 
             except SQLAlchemyError as e:
-                self.session.rollback()
                 logging.error(f"❌ Database error while saving detection: {e}")

nedo_vision_worker_core/repositories/WorkerSourcePipelineDebugRepository.py

@@ -3,15 +3,14 @@ import json
 import os
 import cv2
 import numpy as np
-from sqlalchemy.orm import Session
+from .BaseRepository import BaseRepository
 from ..database.DatabaseManager import DatabaseManager
 from ..models.worker_source_pipeline_debug import WorkerSourcePipelineDebugEntity
 
 
-class WorkerSourcePipelineDebugRepository:
+class WorkerSourcePipelineDebugRepository(BaseRepository):
     def __init__(self):
-        self.db_manager = DatabaseManager()
-        self.session: Session = self.db_manager.get_session("default")
+        super().__init__(db_name="default")
         self.storage_dir = DatabaseManager.STORAGE_PATHS["files"] / "debug_image"
         os.makedirs(self.storage_dir, exist_ok=True)
 
@@ -21,24 +20,25 @@ class WorkerSourcePipelineDebugRepository:
 
         :return: A list of pipeline IDs (str) with null data.
         """
-        now = datetime.now(timezone.utc)
-        cutoff_time = now - timedelta(minutes=1)
+        with self._get_session() as session:
+            now = datetime.now(timezone.utc)
+            cutoff_time = now - timedelta(minutes=1)
 
-        self.session.query(WorkerSourcePipelineDebugEntity)\
-            .filter(
-                WorkerSourcePipelineDebugEntity.data == None,
-                WorkerSourcePipelineDebugEntity.created_at < cutoff_time
-            ).delete(synchronize_session=False)
+            session.query(WorkerSourcePipelineDebugEntity)\
+                .filter(
+                    WorkerSourcePipelineDebugEntity.data == None,
+                    WorkerSourcePipelineDebugEntity.created_at < cutoff_time
+                ).delete(synchronize_session=False)
 
-        self.session.commit()
+            session.commit()
 
-        results = self.session.query(
-            WorkerSourcePipelineDebugEntity.worker_source_pipeline_id
-        ).filter(
-            WorkerSourcePipelineDebugEntity.data == None
-        ).distinct().all()
+            results = session.query(
+                WorkerSourcePipelineDebugEntity.worker_source_pipeline_id
+            ).filter(
+                WorkerSourcePipelineDebugEntity.data == None
+            ).distinct().all()
 
-        return [row[0] for row in results]
+            return [row[0] for row in results]
 
     def update_debug_entries_by_pipeline_id(self, pipeline_id: int, image, data: str):
         """
@@ -48,34 +48,35 @@ class WorkerSourcePipelineDebugRepository:
         :param new_data: The new data to update the entries with.
         :return: The number of updated entries.
         """
-        now = datetime.now(timezone.utc)
-        current_datetime = now.strftime("%Y%m%d_%H%M%S")
+        with self._get_session() as session:
+            now = datetime.now(timezone.utc)
+            current_datetime = now.strftime("%Y%m%d_%H%M%S")
 
-        stringified_data = json.dumps(
-            {
-                "timestamp": now.strftime("%Y-%m-%dT%H:%M:%SZ"),
-                "tracked_objects": data,
-            },
-            default=lambda o: (
-                float(o) if isinstance(o, np.floating) else
-                int(o) if isinstance(o, np.integer) else
-                list(o) if isinstance(o, (np.ndarray, tuple)) else
-                str(o)
-            )
-        ) 
-        
-        full_image_filename = f"{pipeline_id}_{current_datetime}.jpg"
-        full_image_path = os.path.join(self.storage_dir, full_image_filename)
-        cv2.imwrite(full_image_path, image)
-
-        updated_entries = self.session.query(WorkerSourcePipelineDebugEntity)\
-            .filter_by(worker_source_pipeline_id=pipeline_id)\
-            .update(
+            stringified_data = json.dumps(
                 {
-                    "image_path": full_image_path,
-                    "data": stringified_data
+                    "timestamp": now.strftime("%Y-%m-%dT%H:%M:%SZ"),
+                    "tracked_objects": data,
                 },
-                synchronize_session="fetch"
-            )
-        self.session.commit()
-        return updated_entries
+                default=lambda o: (
+                    float(o) if isinstance(o, np.floating) else
+                    int(o) if isinstance(o, np.integer) else
+                    list(o) if isinstance(o, (np.ndarray, tuple)) else
+                    str(o)
+                )
+            ) 
+            
+            full_image_filename = f"{pipeline_id}_{current_datetime}.jpg"
+            full_image_path = os.path.join(self.storage_dir, full_image_filename)
+            cv2.imwrite(full_image_path, image)
+
+            updated_entries = session.query(WorkerSourcePipelineDebugEntity)\
+                .filter_by(worker_source_pipeline_id=pipeline_id)\
+                .update(
+                    {
+                        "image_path": full_image_path,
+                        "data": stringified_data
+                    },
+                    synchronize_session="fetch"
+                )
+            # Commit happens automatically via context manager
+            return updated_entries

nedo_vision_worker_core/repositories/WorkerSourcePipelineDetectionRepository.py

@@ -3,16 +3,15 @@ import json
 import os
 import cv2
 import numpy as np
-from sqlalchemy.orm import Session
+from .BaseRepository import BaseRepository
 from ..ai.FrameDrawer import FrameDrawer
 from ..database.DatabaseManager import DatabaseManager
 from ..models.worker_source_pipeline_detection import WorkerSourcePipelineDetectionEntity
 
 
-class WorkerSourcePipelineDetectionRepository:
+class WorkerSourcePipelineDetectionRepository(BaseRepository):
     def __init__(self):
-        self.db_manager = DatabaseManager()
-        self.session: Session = self.db_manager.get_session("default")
+        super().__init__(db_name="default")
         self.storage_dir = DatabaseManager.STORAGE_PATHS["files"] / "detection_image"
         os.makedirs(self.storage_dir, exist_ok=True)
 
@@ -56,16 +55,16 @@ class WorkerSourcePipelineDetectionRepository:
         ) 
 
         try:
-            new_detection = WorkerSourcePipelineDetectionEntity(
-                worker_source_pipeline_id=pipeline_id,
-                image_path=full_image_path,
-                data=stringified_data,
-                created_at=datetime.utcnow()
-            )
-            self.session.add(new_detection)
-            self.session.flush()
-            self.session.commit()
-            print(f"✅ Inserted detection data for pipeline {pipeline_id}")
+            with self._get_session() as session:
+                new_detection = WorkerSourcePipelineDetectionEntity(
+                    worker_source_pipeline_id=pipeline_id,
+                    image_path=full_image_path,
+                    data=stringified_data,
+                    created_at=datetime.utcnow()
+                )
+                session.add(new_detection)
+                session.flush()
+                # Commit happens automatically via context manager
+                print(f"✅ Inserted detection data for pipeline {pipeline_id}")
         except Exception as e:
-            self.session.rollback()
             print(f"❌ Database error while saving detection: {e}")

nedo_vision_worker_core/repositories/WorkerSourcePipelineRepository.py

@@ -1,15 +1,13 @@
 import json
-from sqlalchemy.orm import Session
 from sqlalchemy.exc import SQLAlchemyError
-from ..database.DatabaseManager import DatabaseManager
+from .BaseRepository import BaseRepository
 from ..models.worker_source_pipeline import WorkerSourcePipelineEntity
 from ..models.worker_source_pipeline_config import WorkerSourcePipelineConfigEntity
 
 
-class WorkerSourcePipelineRepository:
+class WorkerSourcePipelineRepository(BaseRepository):
     def __init__(self):
-        self.db_manager = DatabaseManager()
-        self.session: Session = self.db_manager.get_session("config")
+        super().__init__(db_name="config")
 
     def get_all_pipelines(self):
         """
@@ -18,8 +16,14 @@ class WorkerSourcePipelineRepository:
         Returns:
             list: A list of WorkerSourcePipelineEntity records.
         """
-        self.session.expire_all()
-        return self.session.query(WorkerSourcePipelineEntity).all()
+        with self._get_session() as session:
+            session.expire_all()
+            # Query and detach from session before returning
+            pipelines = session.query(WorkerSourcePipelineEntity).all()
+            # Expunge objects so they can be used outside session
+            for pipeline in pipelines:
+                session.expunge(pipeline)
+            return pipelines
 
     def get_pipeline_configs_by_pipeline_id(self, pipeline_id):
         """
@@ -38,42 +42,70 @@ class WorkerSourcePipelineRepository:
             dict: A dictionary mapping pipeline_config_code to its configuration details.
         """
         try:
-            pipeline_configs = (
-                self.session.query(WorkerSourcePipelineConfigEntity)
-                .filter(WorkerSourcePipelineConfigEntity.worker_source_pipeline_id == pipeline_id)
-                .all()
-            )
+            with self._get_session() as session:
+                pipeline_configs = (
+                    session.query(WorkerSourcePipelineConfigEntity)
+                    .filter(WorkerSourcePipelineConfigEntity.worker_source_pipeline_id == pipeline_id)
+                    .all()
+                )
 
-            def parse_value(value):
-                """Attempts to parse the value as JSON if applicable."""
-                if not value:
-                    return value  # Keep None or empty string as is
-                
-                value = value.strip()  # Remove leading/trailing spaces
-                if (value.startswith("{") and value.endswith("}")) or (value.startswith("[") and value.endswith("]")):
-                    try:
-                        return json.loads(value)  # Parse JSON object or list
-                    except json.JSONDecodeError:
-                        pass  # Keep as string if parsing fails
-                return value  # Return original value if not JSON
+                def parse_value(value):
+                    """Attempts to parse the value as JSON if applicable."""
+                    if not value:
+                        return value  # Keep None or empty string as is
+                    
+                    value = value.strip()  # Remove leading/trailing spaces
+                    if (value.startswith("{") and value.endswith("}")) or (value.startswith("[") and value.endswith("]")):
+                        try:
+                            return json.loads(value)  # Parse JSON object or list
+                        except json.JSONDecodeError:
+                            pass  # Keep as string if parsing fails
+                    return value  # Return original value if not JSON
 
-            # Convert result into a dictionary with pipeline_config_code as key
-            config_dict = {
-                config.pipeline_config_code: {
-                    "id": config.id,
-                    "is_enabled": config.is_enabled,  # Keep original boolean value
-                    "value": parse_value(config.value),  # Parse JSON if applicable
-                    "name": config.pipeline_config_name
+                # Convert result into a dictionary with pipeline_config_code as key
+                config_dict = {
+                    config.pipeline_config_code: {
+                        "id": config.id,
+                        "is_enabled": config.is_enabled,  # Keep original boolean value
+                        "value": parse_value(config.value),  # Parse JSON if applicable
+                        "name": config.pipeline_config_name
+                    }
+                    for config in pipeline_configs
                 }
-                for config in pipeline_configs
-            }
 
-            return config_dict
+                return config_dict
 
         except SQLAlchemyError as e:
             print(f"Database error while retrieving pipeline configs: {e}")
             return {}
         
     def get_worker_source_pipeline(self, pipeline_id):
-        self.session.expire_all()
-        return self.session.query(WorkerSourcePipelineEntity).filter_by(id=pipeline_id).first()
+        with self._get_session() as session:
+            session.expire_all()
+            pipeline = session.query(WorkerSourcePipelineEntity).filter_by(id=pipeline_id).first()
+            if pipeline:
+                session.expunge(pipeline)  # Detach from session
+            return pipeline
+    
+    def update_pipeline_status(self, pipeline_id: str, status_code: str) -> bool:
+        """
+        Update the status of a pipeline.
+        
+        Args:
+            pipeline_id: The ID of the pipeline to update
+            status_code: The new status code ('run', 'stop', 'restart')
+            
+        Returns:
+            bool: True if update was successful, False otherwise
+        """
+        try:
+            with self._get_session() as session:
+                pipeline = session.query(WorkerSourcePipelineEntity).filter_by(id=pipeline_id).first()
+                if pipeline:
+                    pipeline.pipeline_status_code = status_code
+                    # Commit happens automatically via context manager
+                    return True
+                return False
+        except SQLAlchemyError as e:
+            print(f"Database error while updating pipeline status: {e}")
+            return False

nedo_vision_worker_core/repositories/WorkerSourceRepository.py

@@ -1,12 +1,10 @@
-from sqlalchemy.orm import Session
-from ..database.DatabaseManager import DatabaseManager
+from .BaseRepository import BaseRepository
 from ..models.worker_source import WorkerSourceEntity
 
 
-class WorkerSourceRepository:
+class WorkerSourceRepository(BaseRepository):
     def __init__(self):
-        self.db_manager = DatabaseManager()
-        self.session: Session = self.db_manager.get_session("config")
+        super().__init__(db_name="config")
 
     def get_worker_sources(self):
         """
@@ -15,5 +13,9 @@ class WorkerSourceRepository:
         Returns:
             list: A list of WorkerSourceEntity records.
         """
-        self.session.expire_all()
-        return self.session.query(WorkerSourceEntity).all()
+        with self._get_session() as session:
+            session.expire_all()
+            sources = session.query(WorkerSourceEntity).all()
+            for source in sources:
+                session.expunge(source)
+            return sources