gst-python-ml 0.1.0__py3-none-any.whl

engine/gst_device_queue_pool.py

@@ -0,0 +1,121 @@
+# GstDeviceQueuePool
+# Copyright (C) 2024-2025 Collabora Ltd.
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Library General Public
+# License as published by the Free Software Foundation; either
+# version 2 of the License, or (at your option) any later version.
+#
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Library General Public License for more details.
+#
+# You should have received a copy of the GNU Library General Public
+# License along with this library; if not, write to the
+# Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
+# Boston, MA 02110-1301, USA.
+
+import gi
+from abc import ABC, abstractmethod
+
+gi.require_version("Gst", "1.0")
+from gi.repository import Gst  # noqa: E402
+
+
+class GstDeviceQueue(ABC):
+    def __init__(self, queue_handle):
+        """
+        Initialize the DeviceQueue with a given queue handle.
+
+        :param queue_handle: An integer representing the queue handle.
+        """
+        self.queue_handle = queue_handle
+
+    @abstractmethod
+    def synchronize(self):
+        """
+        Abstract method to synchronize the device queue.
+        Must be implemented by subclasses.
+        """
+        pass
+
+    def __repr__(self):
+        return f"DeviceQueue(handle={self.queue_handle})"
+
+
+class GstDeviceQueuePool:
+    def __init__(self):
+        """
+        Initialize the DeviceQueuePool with an empty dictionary to map IDs to DeviceQueues.
+        """
+        self.queues = {}
+
+    def add_queue(self, queue_id, device_queue):
+        """
+        Add a DeviceQueue to the pool by its ID.
+
+        :param queue_id: Unique ID for the DeviceQueue.
+        :param device_queue: A DeviceQueue object to add to the pool.
+        """
+        if queue_id in self.queues:
+            Gst.warning(
+                f"DeviceQueue with ID {queue_id} already exists. Not adding again."
+            )
+            return
+
+        self.queues[queue_id] = device_queue
+
+    def get_queue(self, queue_id):
+        """
+        Retrieve a DeviceQueue by its ID.
+
+        :param queue_id: Unique ID of the queue in the pool.
+        :return: DeviceQueue object if found, None otherwise.
+        """
+        queue = self.queues.get(queue_id, None)
+        if queue is None:
+            Gst.warning(f"No DeviceQueue found for ID {queue_id}.")
+        return queue
+
+    def __repr__(self):
+        return f"GstDeviceQueuePool(queues={self.queues})"
+
+
+class DeviceQueueManager:
+    _instance = None
+    _device_queue_pools = {}
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(DeviceQueueManager, cls).__new__(cls)
+            cls._instance._device_queue_pools = {}
+        return cls._instance
+
+    def add_pool(self, device, queue_pool):
+        """
+        Adds a DeviceQueuePool for a specific device.
+
+        :param device: Unique identifier for the device (e.g., "cuda:0").
+        :param queue_pool: A DeviceQueuePool object to associate with the device.
+        """
+        if device in self._device_queue_pools:
+            return  # Do not add if it already exists
+
+        self._device_queue_pools[device] = queue_pool
+        Gst.info(f"Added DeviceQueuePool for device {device}.")
+
+    def get_pool(self, device):
+        """
+        Retrieves the DeviceQueuePool associated with the specified device.
+
+        :param device: Unique identifier for the device.
+        :return: DeviceQueuePool object if found, None otherwise.
+        """
+        pool = self._device_queue_pools.get(device, None)
+        if pool is None:
+            Gst.warning(f"No DeviceQueuePool found for device {device}.")
+        return pool
+
+    def __repr__(self):
+        return f"DeviceQueueManager(device_queue_pools={self._device_queue_pools})"

engine/gst_engine_factory.py

@@ -0,0 +1,130 @@
+# GstEngineFactory
+# Copyright (C) 2024-2025 Collabora Ltd.
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Library General Public
+# License as published by the Free Software Foundation; either
+# version 2 of the License, or (at your option) any later version.
+#
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Library General Public License for more details.
+#
+# You should have received a copy of the GNU Library General Public
+# License along with this library; if not, write to the
+# Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
+# Boston, MA 02110-1301, USA.
+
+
+try:
+    from .gst_pytorch_engine import GstPyTorchEngine
+
+    _pytorch_engine_available = True
+except ImportError:
+    _pytorch_engine_available = False
+
+try:
+    from .gst_pytorch_yolo_engine import GstPyTorchYoloEngine
+
+    _pytorch_yolo_engine_available = True
+except ImportError:
+    _pytorch_yolo_engine_available = False
+
+try:
+    from .gst_tflite_engine import GstTFLiteEngine
+
+    _tflite_engine_available = True
+except ImportError:
+    _tflite_engine_available = False
+
+try:
+    from .gst_tensorflow_engine import GstTensorFlowEngine
+
+    _tensorflow_engine_available = True
+except ImportError:
+    _tensorflow_engine_available = False
+
+try:
+    from .gst_onnx_engine import GstONNXEngine
+
+    _onnx_engine_available = True
+except ImportError:
+    _onnx_engine_available = False
+
+try:
+    from .gst_openvino_engine import GstOpenVinoEngine
+
+    _openvino_engine_available = True
+except ImportError:
+    _openvino_engine_available = False
+
+
+class GstEngineFactory:
+    # Define the constant strings for each engine
+    PYTORCH_ENGINE = "pytorch"
+    PYTORCH_YOLO_ENGINE = "pytorch-yolo"
+    TFLITE_ENGINE = "tflite"
+    TENSORFLOW_ENGINE = "tensorflow"
+    ONNX_ENGINE = "onnx"
+    OPENVINO_ENGINE = "openvino"
+
+    @staticmethod
+    def create_engine(engine_type, device="cpu"):
+        """
+        Factory method to create the appropriate engine based on the engine type.
+
+        :param engine_type: The type of the ML engine, e.g., "pytorch" or "tflite".
+        :param device: The device to run the engine on (default is "cpu").
+        :return: An instance of the appropriate ML engine class.
+        """
+        if engine_type == GstEngineFactory.PYTORCH_ENGINE:
+            if _pytorch_engine_available:
+                return GstPyTorchEngine(device)
+            else:
+                raise ImportError(
+                    f"{GstEngineFactory.PYTORCH_ENGINE} engine is not available."
+                )
+
+        if engine_type == GstEngineFactory.PYTORCH_YOLO_ENGINE:
+            if _pytorch_yolo_engine_available:
+                return GstPyTorchYoloEngine(device)
+            else:
+                raise ImportError(
+                    f"{GstEngineFactory.PYTORCH_YOLO_ENGINE} engine is not available."
+                )
+
+        elif engine_type == GstEngineFactory.TFLITE_ENGINE:
+            if _tflite_engine_available:
+                return GstTFLiteEngine(device)
+            else:
+                raise ImportError(
+                    f"{GstEngineFactory.TFLITE_ENGINE} engine is not available."
+                )
+
+        elif engine_type == GstEngineFactory.TENSORFLOW_ENGINE:
+            if _tensorflow_engine_available:
+                return GstTensorFlowEngine(device)
+            else:
+                raise ImportError(
+                    f"{GstEngineFactory.TENSORFLOW_ENGINE} engine is not available."
+                )
+
+        elif engine_type == GstEngineFactory.ONNX_ENGINE:
+            if _onnx_engine_available:
+                return GstONNXEngine(device)
+            else:
+                raise ImportError(
+                    f"{GstEngineFactory.ONNX_ENGINE} engine is not available."
+                )
+
+        elif engine_type == GstEngineFactory.OPENVINO_ENGINE:
+            if _openvino_engine_available:
+                return GstOpenVinoEngine(device)
+            else:
+                raise ImportError(
+                    f"{GstEngineFactory.OPENVINO_ENGINE} engine is not available."
+                )
+
+        else:
+            raise ValueError(f"Unsupported engine type: {engine_type}")

engine/gst_ml_engine.py

@@ -0,0 +1,75 @@
+# GstMLEngine
+# Copyright (C) 2024-2025 Collabora Ltd.
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Library General Public
+# License as published by the Free Software Foundation; either
+# version 2 of the License, or (at your option) any later version.
+#
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Library General Public License for more details.
+#
+# You should have received a copy of the GNU Library General Public
+# License along with this library; if not, write to the
+# Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
+# Boston, MA 02110-1301, USA.
+
+from abc import ABC, abstractmethod
+
+
+class GstMLEngine(ABC):
+    def __init__(self, device="cpu"):
+        self.device = device
+        self.device_index = 0
+        self.model = None
+        self.vision_language_model = False
+        self.tokenizer = None
+        self.image_processor = None
+        self.batch_size = 1
+        self.frame_buffer = []
+        self.counter = 0
+        self.device_queue_id = None
+        self.track = False
+        self.prompt = "What is shown in this image?"  # Default prompt
+
+    @abstractmethod
+    def load_model(self, model_name, **kwargs):
+        """Load a model by name or path, with additional options."""
+        pass
+
+    def set_prompt(self, prompt):
+        """Set the custom prompt for generating responses."""
+        self.prompt = prompt
+
+    def get_prompt(self):
+        """Return the custom prompt."""
+        return self.prompt
+
+    def get_device(self):
+        """Return the loaded model for use in inference."""
+        return self.device
+
+    @abstractmethod
+    def set_device(self, device):
+        """Set the device (e.g., cpu, cuda)."""
+        pass
+
+    def get_model(self):
+        """Return the loaded model for use in inference."""
+        return self.model
+
+    def set_model(self, model):
+        """Set the model directly (useful for loading pre-built models)."""
+        self.model = model
+
+    @abstractmethod
+    def forward(self, frame):
+        """Execute inference (usually object detection)"""
+        pass
+
+    @abstractmethod
+    def generate(self, input_text, max_length=100):
+        """Generate LLM text"""
+        pass

engine/gst_onnx_engine.py

@@ -0,0 +1,99 @@
+# GstONNXEngine
+# Copyright (C) 2024-2025 Collabora Ltd.
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Library General Public
+# License as published by the Free Software Foundation; either
+# version 2 of the License, or (at your option) any later version.
+#
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Library General Public License for more details.
+#
+# You should have received a copy of the GNU Library General Public
+# License along with this library; if not, write to the
+# Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
+# Boston, MA 02110-1301, USA.
+
+import gi
+import numpy as np
+import onnxruntime as ort  # ONNX Runtime for executing ONNX models
+from .gst_ml_engine import GstMLEngine
+
+gi.require_version("Gst", "1.0")
+gi.require_version("GstBase", "1.0")
+gi.require_version("GLib", "2.0")
+from gi.repository import Gst  # noqa: E402
+
+
+class GstONNXEngine(GstMLEngine):
+    def __init__(self, device="cpu"):
+        """
+        Initialize the ONNX engine with the specified device.
+        """
+        super().__init__(device)
+        self.session = None
+        self.input_names = None
+        self.output_names = None
+
+    def load_model(self, model_name, **kwargs):
+        """
+        Load the ONNX model from the specified file path.
+        """
+        try:
+            # Create the ONNX runtime session with device (CPU, CUDA, etc.)
+            providers = (
+                ["CPUExecutionProvider"]
+                if self.device == "cpu"
+                else ["CUDAExecutionProvider"]
+            )
+            self.session = ort.InferenceSession(model_name, providers=providers)
+
+            # Extract input and output names for reference
+            self.input_names = [inp.name for inp in self.session.get_inputs()]
+            self.output_names = [out.name for out in self.session.get_outputs()]
+
+            Gst.info(f"ONNX model '{model_name}' loaded successfully on {self.device}.")
+        except Exception as e:
+            Gst.error(f"Failed to load ONNX model '{model_name}'. Error: {e}")
+
+    def set_device(self, device):
+        """
+        Set the device for inference.
+        """
+        self.device = device
+        # ONNX Runtime does not allow changing the device after session creation,
+        # so we need to reload the model if device changes.
+        if self.session:
+            model_path = (
+                self.session.get_modelmeta().producer_name
+            )  # Assuming model path is stored
+            self.load_model(model_path)
+
+    def forward(self, frame):
+        """
+        Perform inference on the given frame using the ONNX model.
+        """
+        try:
+            # Preprocess the frame (resize, normalize, etc.) as required by the model
+            input_tensor = np.expand_dims(
+                frame.astype(np.float32), axis=0
+            )  # Add batch dimension
+
+            # Prepare the input dictionary
+            input_dict = {self.input_names[0]: input_tensor}
+
+            # Run inference
+            output_data = self.session.run(self.output_names, input_dict)
+
+            # Convert output to NumPy arrays if needed
+            results = [np.array(out) for out in output_data]
+
+            return (
+                results if len(results) > 1 else results[0]
+            )  # Return single result or list of results
+
+        except Exception as e:
+            Gst.error(f"Error during ONNX inference: {e}")
+            return None

engine/gst_openvino_engine.py

@@ -0,0 +1,150 @@
+# GstOpenVinoEngine
+# Copyright (C) 2024-2025 Collabora Ltd.
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Library General Public
+# License as published by the Free Software Foundation; either
+# version 2 of the License, or (at your option) any later version.
+#
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Library General Public License for more details.
+#
+# You should have received a copy of the GNU Library General Public
+# License along with this library; if not, write to the
+# Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
+# Boston, MA 02110-1301, USA.
+
+import gi
+import numpy as np
+from openvino.runtime import Core
+from .gst_ml_engine import GstMLEngine
+
+gi.require_version("Gst", "1.0")
+from gi.repository import Gst  # noqa: E402
+
+
+class GstOpenVinoEngine(GstMLEngine):
+    def __init__(self, device="CPU"):
+        self.device = device
+        self.core = Core()
+        self.compiled_model = None
+        self.tokenizer = None
+        self.is_vision_model = False
+        self.is_llm = False
+
+    def load_model(self, model_name, **kwargs):
+        """
+        Load the model using OpenVINO and differentiate between vision and LLM models.
+        """
+        try:
+            # Load model
+            model_path = f"{model_name}.xml"
+            self.model = self.core.read_model(model=model_path)
+            self.compiled_model = self.core.compile_model(self.model, self.device)
+            Gst.info(f"Model '{model_name}' loaded successfully on {self.device}")
+
+            # Inspect input shape to determine the type of model
+            input_shape = self.compiled_model.input(0).shape
+
+            if (
+                len(input_shape) == 4 and input_shape[1] == 3
+            ):  # Expecting image input (batch, 3 channels, height, width)
+                self.is_vision_model = True
+                Gst.info("Model identified as a vision model.")
+            elif len(input_shape) == 2:  # Expecting LLM input (batch, sequence_length)
+                self.is_llm = True
+                Gst.info("Model identified as a large language model (LLM).")
+
+            if self.is_llm:
+                # Load a tokenizer for the LLM (Hugging Face for instance)
+                from transformers import AutoTokenizer
+
+                self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+        except Exception as e:
+            raise ValueError(f"Failed to load model '{model_name}'. Error: {e}")
+
+    def generate(self, input_data, max_length=100):
+        """
+        Process input for the respective model type.
+        """
+        if self.is_vision_model:
+            # Preprocess input for vision models (e.g., resize, normalize)
+            return self.preprocess_vision_input(input_data)
+        elif self.is_llm:
+            # Preprocess input for LLMs (tokenization)
+            return self.preprocess_llm_input(input_data, max_length)
+        else:
+            raise ValueError("Unknown model type. Please load a model first.")
+
+    def preprocess_vision_input(self, image):
+        """
+        Preprocess input image for vision models.
+        """
+        # Example preprocessing: Resize and normalize the image
+        input_shape = self.compiled_model.input(0).shape
+        resized_image = np.resize(
+            image, input_shape[2:]
+        )  # Resize image to (height, width)
+        resized_image = resized_image.astype(np.float32) / 255.0  # Normalize
+        resized_image = np.expand_dims(resized_image, axis=0)  # Add batch dimension
+        return resized_image
+
+    def preprocess_llm_input(self, input_text, max_length):
+        """
+        Tokenize input text for LLM models.
+        """
+        if self.tokenizer is None:
+            raise RuntimeError("Tokenizer not initialized. Load an LLM model first.")
+
+        tokens = self.tokenizer(
+            input_text,
+            max_length=max_length,
+            truncation=True,
+            padding="max_length",
+            return_tensors="np",  # Return as NumPy array
+        )
+        input_ids = tokens["input_ids"]
+        attention_mask = tokens["attention_mask"]
+
+        return input_ids, attention_mask
+
+    def forward(self, input_data):
+        """
+        Perform inference using the loaded OpenVINO model.
+        """
+        if self.compiled_model is None:
+            raise RuntimeError("Model not loaded. Please load the model first.")
+
+        if self.is_vision_model:
+            # Perform inference for vision models
+            return self.perform_vision_inference(input_data)
+        elif self.is_llm:
+            # Perform inference for LLM models
+            return self.perform_llm_inference(input_data)
+
+    def perform_vision_inference(self, input_data):
+        """
+        Perform inference on vision models.
+        """
+        infer_request = self.compiled_model.create_infer_request()
+        infer_request.infer({self.compiled_model.input(0): input_data})
+        return infer_request.get_output_tensor(self.compiled_model.output(0)).data
+
+    def perform_llm_inference(self, input_data):
+        """
+        Perform inference on LLM models.
+        """
+        input_ids, attention_mask = (
+            input_data  # Assuming input_data is (input_ids, attention_mask)
+        )
+        infer_request = self.compiled_model.create_infer_request()
+        infer_request.infer(
+            {
+                self.compiled_model.input(0): input_ids,
+                self.compiled_model.input(1): attention_mask,
+            }
+        )
+        return infer_request.get_output_tensor(self.compiled_model.output(0)).data