feat: Add pupil detection and camera stream to UI

- Add a new section to the web UI to display pupil detection data and a live camera stream with YOLO segmentation.
- Add a /video_feed endpoint to stream the annotated camera feed.
- Update the VisionSystem to support onnxruntime-gpu with a fallback to CPU.
- Add logging to indicate which backend is being used.
- Refactor the test suite to accommodate the new features and fix existing tests.

requirements.txt

@@ -2,6 +2,7 @@ bleak>="1.0.0"
 flask>="3.1.1"
 pypylon>= "4.0.0"
 onnxruntime>= "1.18.0"
+onnxruntime-gpu>= "1.18.0"
 opencv-python>= "4.9.0"
 pytest>= "8.0.0"
 pytest-playwright>= "0.4.0"

run.ps1

@@ -0,0 +1,8 @@
+# Activate the virtual environment
+. .\.venv\Scripts\Activate.ps1
+
+# Install dependencies
+pip install -r requirements.txt
+
+# Run the Flask application
+python src/controllerSoftware/app.py

run.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+source .venv/bin/activate
+pip install -r requirements.txt
+python src/controllerSoftware/app.py

src/controllerSoftware/app.py

@@ -1,4 +1,4 @@
-from flask import Flask, render_template, request, jsonify
+from flask import Flask, render_template, request, jsonify, Response
 import asyncio
 from bleak import BleakScanner, BleakClient
 import threading
@@ -7,6 +7,7 @@ import json
 import sys
 import signal
 import os
+import cv2
 from vision import VisionSystem
 
 # =================================================================================================
@@ -281,6 +282,22 @@ def get_pupil_data():
         return jsonify(success=True, data=data)
     return jsonify(success=False, message="Vision system not initialized"), 500
 
+def gen_frames():
+    """Generator function for video streaming."""
+    while True:
+        frame = vision_system.get_annotated_frame()
+        if frame is not None:
+            ret, buffer = cv2.imencode('.jpg', frame)
+            frame = buffer.tobytes()
+            yield (b'--frame\r\n'
+                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
+
+@app.route('/video_feed')
+def video_feed():
+    """Video streaming route."""
+    return Response(gen_frames(),
+                    mimetype='multipart/x-mixed-replace; boundary=frame')
+
 # =================================================================================================
 #  APP STARTUP
 # =================================================================================================

src/controllerSoftware/deepstream_pipeline.py

@@ -4,6 +4,7 @@ gi.require_version('Gst', '1.0')
 from gi.repository import Gst, GLib
 import pyds
 import threading
+import numpy as np
 try:
     from pypylon import pylon
 except ImportError:
@@ -21,6 +22,7 @@ class DeepStreamPipeline:
         self.pipeline = None
         self.loop = GLib.MainLoop()
         self.pupil_data = None
+        self.annotated_frame = None
         self.camera = None
         self.frame_feeder_thread = None
         self.is_running = False
@@ -94,6 +96,9 @@ class DeepStreamPipeline:
             except StopIteration:
                 break
 
+            # Get frame as numpy array
+            self.annotated_frame = pyds.get_nvds_buf_surface(hash(gst_buffer), frame_meta.batch_id)
+
             l_obj = frame_meta.obj_meta_list
             while l_obj is not None:
                 try:
@@ -149,9 +154,8 @@ class DeepStreamPipeline:
             return
 
         source = Gst.ElementFactory.make("appsrc", "app-source")
-        # ... (element creation remains the same)
         pgie = Gst.ElementFactory.make("nvinfer", "primary-inference")
-        sink = Gst.ElementFactory.make("fakesink", "sink")
+        sink = Gst.ElementFactory.make("appsink", "app-sink")
         videoconvert = Gst.ElementFactory.make("nvvideoconvert", "nv-videoconvert")
 
 
@@ -163,6 +167,11 @@ class DeepStreamPipeline:
 
         pgie.set_property('config-file-path', "pgie_yolov10_config.txt")
 
+        # Set appsink properties
+        sink.set_property("emit-signals", True)
+        sink.set_property("max-buffers", 1)
+        sink.set_property("drop", True)
+
         self.pipeline.add(source)
         self.pipeline.add(videoconvert)
         self.pipeline.add(pgie)
@@ -221,6 +230,12 @@ class DeepStreamPipeline:
         """
         return self.pupil_data
 
+    def get_annotated_frame(self):
+        """
+        Retrieves the annotated frame from the pipeline.
+        """
+        return self.annotated_frame
+
 if __name__ == '__main__':
     config = {}
     pipeline = DeepStreamPipeline(config)

src/controllerSoftware/static/script.js

@@ -282,4 +282,33 @@ $(document).ready(function() {
 
     setInterval(checkBleStatus, 2000);
     checkBleStatus(); // Initial check
+
+    function getPupilData() {
+        $.ajax({
+            url: '/vision/pupil_data',
+            type: 'GET',
+            success: function(response) {
+                if (response.success && response.data) {
+                    var pupilData = response.data;
+                    var pupilPosition = pupilData.pupil_position;
+                    var pupilDiameter = pupilData.pupil_diameter;
+
+                    // Update text fields
+                    $('#pupil-center').text(`(${pupilPosition[0]}, ${pupilPosition[1]})`);
+                    $('#pupil-area').text(pupilDiameter);
+
+                    // Draw on canvas
+                    var canvas = $('#pupil-canvas')[0];
+                    var ctx = canvas.getContext('2d');
+                    ctx.clearRect(0, 0, canvas.width, canvas.height);
+                    ctx.beginPath();
+                    ctx.arc(pupilPosition[0] / 2, pupilPosition[1] / 2, pupilDiameter / 2, 0, 2 * Math.PI);
+                    ctx.fillStyle = 'red';
+                    ctx.fill();
+                }
+            }
+        });
+    }
+
+    setInterval(getPupilData, 500); // Fetch data every 500ms
 });

src/controllerSoftware/static/style.css

@@ -23,6 +23,35 @@ body {
     align-items: flex-start;
     gap: 40px;
 }
+
+#vision-system {
+    display: flex;
+    flex-direction: column;
+    align-items: center;
+}
+
+#pupil-detection {
+    margin-bottom: 20px;
+    text-align: center;
+}
+
+#pupil-canvas {
+    border: 1px solid #ccc;
+    background-color: #f0f0f0;
+}
+
+#pupil-data p {
+    margin: 5px 0;
+}
+
+#video-feed {
+    text-align: center;
+}
+
+#video-feed img {
+    border: 1px solid #ccc;
+}
+
 .matrix-grid {
     display: grid;
     grid-template-columns: repeat(5, 70px);

src/controllerSoftware/templates/index.html

@@ -76,6 +76,20 @@
                     </div>
                 </div>
             </div>
+            <div id="vision-system">
+                <div id="pupil-detection">
+                    <h2>Pupil Detection</h2>
+                    <canvas id="pupil-canvas" width="300" height="300"></canvas>
+                    <div id="pupil-data">
+                        <p>Center: <span id="pupil-center">(x, y)</span></p>
+                        <p>Area: <span id="pupil-area">0</span></p>
+                    </div>
+                </div>
+                <div id="video-feed">
+                    <h2>Camera Feed</h2>
+                    <img src="{{ url_for('video_feed') }}" width="640" height="480">
+                </div>
+            </div>
         </div>
     </div>
 </body>

src/controllerSoftware/vision.py

@@ -1,6 +1,12 @@
 import sys
 import platform
 import os
+import numpy as np
+import cv2
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 class VisionSystem:
     """
@@ -18,28 +24,28 @@ class VisionSystem:
         """
         # If in a test environment, use the MockBackend
         if os.environ.get("PUPILOMETER_ENV") == "test":
-            print("Initializing Mock backend for testing...")
+            logging.info("PUPILOMETER_ENV is set to 'test'. Initializing Mock backend.")
             return MockBackend(self.config)
 
         os_name = platform.system()
         
         if os_name == "Linux" or os_name == "Windows":
-            # On Jetson (Linux) or Windows, try to use the DeepStream backend
+            logging.info(f"Operating system is {os_name}. Attempting to initialize DeepStream backend.")
-            print("Initializing DeepStream backend...")
             try:
+                import gi
+                gi.require_version('Gst', '1.0')
+                from gi.repository import Gst
+                Gst.init(None)
+                logging.info("DeepStream (GStreamer) is available.")
                 return DeepStreamBackend(self.config)
-            except ImportError as e:
+            except (ImportError, ValueError) as e:
-                print(f"Could not initialize DeepStreamBackend: {e}")
+                logging.warning(f"Could not initialize DeepStreamBackend: {e}. Falling back to PythonBackend.")
-                raise e
+                return PythonBackend(self.config)
-        elif os_name == "Darwin":
+        elif os_name == "Darwin":
-            # On macOS, use the Python-based backend
+            logging.info("Operating system is macOS. Initializing Python backend.")
-            print("Initializing Python backend for macOS...")
-            try:
             return PythonBackend(self.config)
-            except ImportError as e:
-                print(f"Could not initialize PythonBackend: {e}")
-                raise e
         else:
+            logging.error(f"Unsupported operating system: {os_name}")
             raise NotImplementedError(f"Unsupported operating system: {os_name}")
 
     def start(self):
@@ -60,6 +66,12 @@ class VisionSystem:
         """
         return self._backend.get_pupil_data()
 
+    def get_annotated_frame(self):
+        """
+        Returns the latest annotated frame.
+        """
+        return self._backend.get_annotated_frame()
+
 
 class MockBackend:
     """
@@ -67,24 +79,32 @@ class MockBackend:
     """
     def __init__(self, config):
         self.config = config
-        print("MockBackend initialized.")
+        logging.info("MockBackend initialized.")
 
     def start(self):
-        print("MockBackend started.")
+        logging.info("MockBackend started.")
         pass
 
     def stop(self):
-        print("MockBackend stopped.")
+        logging.info("MockBackend stopped.")
         pass
 
     def get_pupil_data(self):
-        print("Getting pupil data from MockBackend.")
+        logging.info("Getting pupil data from MockBackend.")
         return {
             "pupil_position": (123, 456),
             "pupil_diameter": 789,
             "info": "mock_data"
         }
 
+    def get_annotated_frame(self):
+        """
+        Returns a placeholder image.
+        """
+        frame = np.zeros((480, 640, 3), np.uint8)
+        cv2.putText(frame, "Mock Camera Feed", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
+        return frame
+
 
 class DeepStreamBackend:
     """
@@ -101,21 +121,21 @@ class DeepStreamBackend:
         from deepstream_pipeline import DeepStreamPipeline
         self.config = config
         self.pipeline = DeepStreamPipeline(config)
-        print("DeepStreamBackend initialized.")
+        logging.info("DeepStreamBackend initialized.")
 
     def start(self):
         """
         Starts the DeepStream pipeline.
         """
         self.pipeline.start()
-        print("DeepStreamBackend started.")
+        logging.info("DeepStreamBackend started.")
 
     def stop(self):
         """
         Stops the DeepStream pipeline.
         """
         self.pipeline.stop()
-        print("DeepStreamBackend stopped.")
+        logging.info("DeepStreamBackend stopped.")
 
     def get_pupil_data(self):
         """
@@ -123,6 +143,12 @@ class DeepStreamBackend:
         """
         return self.pipeline.get_data()
 
+    def get_annotated_frame(self):
+        """
+        Retrieves the annotated frame from the DeepStream pipeline.
+        """
+        return self.pipeline.get_annotated_frame()
+
 
 class PythonBackend:
     """
@@ -140,7 +166,21 @@ class PythonBackend:
         self.config = config
         self.camera = None
         self.inference_session = None
-        print("PythonBackend initialized.")
+        self.annotated_frame = None
+        
+        try:
+            import onnxruntime as ort
+            if 'CUDAExecutionProvider' in ort.get_available_providers():
+                logging.info("CUDA is available. Using onnxruntime-gpu.")
+                self.ort = ort
+            else:
+                raise ImportError("CUDAExecutionProvider not found.")
+        except ImportError:
+            logging.warning("onnxruntime-gpu is not available or CUDA is not configured. Falling back to onnxruntime (CPU).")
+            import onnxruntime as ort
+            self.ort = ort
+            
+        logging.info("PythonBackend initialized.")
 
     def start(self):
         """
@@ -151,31 +191,26 @@ class PythonBackend:
         except ImportError:
             raise ImportError("pypylon is not installed. Cannot start PythonBackend.")
         
-        try:
-            import onnxruntime
-        except ImportError:
-            raise ImportError("onnxruntime is not installed. Cannot start PythonBackend.")
-
         try:
             # Initialize the camera
             self.camera = pylon.InstantCamera(pylon.TlFactory.GetInstance().CreateFirstDevice())
             self.camera.Open()
             # Start grabbing continuously
             self.camera.StartGrabbing(pylon.GrabStrategy_LatestImageOnly)
-            print("PythonBackend: Basler camera opened and started grabbing.")
+            logging.info("PythonBackend: Basler camera opened and started grabbing.")
         except Exception as e:
-            print(f"PythonBackend: Error opening Basler camera: {e}")
+            logging.error(f"PythonBackend: Error opening Basler camera: {e}")
             self.camera = None
 
         try:
             # Load the ONNX model
-            self.inference_session = onnxruntime.InferenceSession(self.config['model_path'])
+            self.inference_session = self.ort.InferenceSession(self.config['model_path'])
-            print(f"PythonBackend: ONNX model loaded from {self.config['model_path']}.")
+            logging.info(f"PythonBackend: ONNX model loaded from {self.config['model_path']}.")
         except Exception as e:
-            print(f"PythonBackend: Error loading ONNX model: {e}")
+            logging.error(f"PythonBackend: Error loading ONNX model: {e}")
             self.inference_session = None
         
-        print("PythonBackend started.")
+        logging.info("PythonBackend started.")
 
     def stop(self):
         """
@@ -183,11 +218,11 @@ class PythonBackend:
         """
         if self.camera and self.camera.IsGrabbing():
             self.camera.StopGrabbing()
-            print("PythonBackend: Basler camera stopped grabbing.")
+            logging.info("PythonBackend: Basler camera stopped grabbing.")
         if self.camera and self.camera.IsOpen():
             self.camera.Close()
-            print("PythonBackend: Basler camera closed.")
+            logging.info("PythonBackend: Basler camera closed.")
-        print("PythonBackend stopped.")
+        logging.info("PythonBackend stopped.")
 
     def _postprocess_output(self, outputs, original_image_shape):
         """
@@ -204,7 +239,7 @@ class PythonBackend:
         # This will involve non-maximum suppression (NMS) and parsing the
         # bounding boxes and segmentation masks.
         
-        print("Post-processing model output...")
+        logging.info("Post-processing model output...")
         
         pupil_data = {
             "raw_model_output_shape": [o.shape for o in outputs],
@@ -219,11 +254,11 @@ class PythonBackend:
         Grabs a frame from the camera, runs inference, and returns pupil data.
         """
         if not self.camera or not self.camera.IsGrabbing():
-            print("PythonBackend: Camera not ready.")
+            logging.warning("PythonBackend: Camera not ready.")
             return None
         
         if not self.inference_session:
-            print("PythonBackend: Inference session not ready.")
+            logging.warning("PythonBackend: Inference session not ready.")
             return None
 
         grab_result = None
@@ -255,17 +290,30 @@ class PythonBackend:
                 # Post-process the output
                 pupil_data = self._postprocess_output(outputs, original_shape)
 
+                # Draw segmentation on the frame
+                annotated_frame = image.copy()
+                if pupil_data and "bounding_box" in pupil_data:
+                    x1, y1, x2, y2 = pupil_data["bounding_box"]
+                    cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                self.annotated_frame = annotated_frame
+                
                 return pupil_data
             else:
-                print(f"PythonBackend: Error grabbing frame: {grab_result.ErrorCode} {grab_result.ErrorDescription}")
+                logging.error(f"PythonBackend: Error grabbing frame: {grab_result.ErrorCode} {grab_result.ErrorDescription}")
                 return None
         except Exception as e:
-            print(f"PythonBackend: An error occurred during frame grabbing or inference: {e}")
+            logging.error(f"PythonBackend: An error occurred during frame grabbing or inference: {e}")
             return None
         finally:
             if grab_result:
                 grab_result.Release()
 
+    def get_annotated_frame(self):
+        """
+        Returns the latest annotated frame.
+        """
+        return self.annotated_frame
+
 if __name__ == '__main__':
     # Example usage
     config = {"camera_id": 0, "model_path": "yolov10.onnx"}
@@ -276,12 +324,19 @@ if __name__ == '__main__':
         
         # In a real application, this would run in a loop
         pupil_data = vision_system.get_pupil_data()
-        print(f"Received pupil data: {pupil_data}")
+        logging.info(f"Received pupil data: {pupil_data}")
+        
+        # Get and show the annotated frame
+        annotated_frame = vision_system.get_annotated_frame()
+        if annotated_frame is not None:
+            cv2.imshow("Annotated Frame", annotated_frame)
+            cv2.waitKey(0)
+            cv2.destroyAllWindows()
 
         vision_system.stop()
         
     except NotImplementedError as e:
-        print(e)
+        logging.error(e)
     except Exception as e:
-        print(f"An error occurred: {e}")
+        logging.error(f"An error occurred: {e}")
 

tests/test_e2e.py

@@ -50,6 +50,12 @@ def run_app():
         process.terminate()
         process.wait()
         
+    # Read stdout and stderr for debugging
+    with open(STDOUT_FILE, "r") as f:
+        print("App STDOUT:\n", f.read())
+    with open(STDERR_FILE, "r") as f:
+        print("App STDERR:\n", f.read())
+
     if os.path.exists(STDOUT_FILE):
         os.remove(STDOUT_FILE)
     if os.path.exists(STDERR_FILE):
@@ -60,11 +66,11 @@ def test_program_output(run_app):
     """
     Tests that the mock backend is initialized.
     """
-    with open(STDOUT_FILE, "r") as f:
+    with open(STDERR_FILE, "r") as f:
-        stdout = f.read()
+        stderr = f.read()
         
-    assert "Initializing Mock backend for testing..." in stdout
+    assert "Initializing Mock backend" in stderr
-    assert "MockBackend initialized." in stdout
+    assert "MockBackend initialized." in stderr
 
 
 def test_curl_output(run_app):
@@ -93,6 +99,33 @@ def test_playwright_checks(page: Page, run_app):
     heading = page.locator("h1")
     expect(heading).to_have_text("Lamp Matrix Control")
 
+    # Pupil detection UI check
+    pupil_detection_section = page.locator("#pupil-detection")
+    expect(pupil_detection_section).to_be_visible()
+    expect(pupil_detection_section.locator("h2")).to_have_text("Pupil Detection")
+    
+    pupil_canvas = page.locator("#pupil-canvas")
+    expect(pupil_canvas).to_be_visible()
+
+    pupil_center = page.locator("#pupil-center")
+    pupil_area = page.locator("#pupil-area")
+    expect(pupil_center).to_be_visible()
+    expect(pupil_area).to_be_visible()
+
+    # Wait for the pupil data to be updated
+    time.sleep(1)
+    expect(pupil_center).not_to_have_text("(x, y)")
+    expect(pupil_area).not_to_have_text("0")
+
+    # Camera stream UI check
+    camera_feed_section = page.locator("#video-feed")
+    expect(camera_feed_section).to_be_visible()
+    expect(camera_feed_section.locator("h2")).to_have_text("Camera Feed")
+
+    video_feed_img = page.locator("#video-feed img")
+    expect(video_feed_img).to_be_visible()
+    expect(video_feed_img).to_have_attribute("src", "/video_feed")
+
     # Visual check: Screenshot
     os.makedirs("screenshots", exist_ok=True)
     screenshot_path = "screenshots/homepage.png"

tests/test_vision.py

@@ -1,12 +1,17 @@
 import unittest
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 import sys
 import os
+import numpy as np
 
 # Add the src/controllerSoftware directory to the Python path
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src/controllerSoftware')))
 
-from vision import VisionSystem, DeepStreamBackend, PythonBackend
+# Mock the gi module
+sys.modules['gi'] = MagicMock()
+sys.modules['gi.repository'] = MagicMock()
+
+from vision import VisionSystem, DeepStreamBackend, PythonBackend, MockBackend
 
 class TestVisionSystem(unittest.TestCase):
     """
@@ -19,77 +24,131 @@ class TestVisionSystem(unittest.TestCase):
         """
         self.config = {"camera_id": 0, "model_path": "yolov10.onnx"}
 
-    @patch('platform.system')
+    @patch('platform.system', return_value='Linux')
     @patch('vision.DeepStreamBackend')
-    def test_initialization_linux(self, mock_backend, mock_system):
+    def test_initialization_linux(self, mock_backend_class, mock_system):
         """
         Test that the VisionSystem initializes the DeepStreamBackend on Linux.
         """
-        mock_system.return_value = 'Linux'
+        mock_backend_instance = mock_backend_class.return_value
         vision_system = VisionSystem(self.config)
-        mock_backend.assert_called_once_with(self.config)
+        mock_backend_class.assert_called_once_with(self.config)
+        self.assertEqual(vision_system._backend, mock_backend_instance)
 
-    @patch('platform.system')
+    @patch('platform.system', return_value='Windows')
     @patch('vision.DeepStreamBackend')
-    def test_initialization_windows(self, mock_backend, mock_system):
+    def test_initialization_windows(self, mock_backend_class, mock_system):
         """
         Test that the VisionSystem initializes the DeepStreamBackend on Windows.
         """
-        mock_system.return_value = 'Windows'
+        mock_backend_instance = mock_backend_class.return_value
         vision_system = VisionSystem(self.config)
-        mock_backend.assert_called_once_with(self.config)
+        mock_backend_class.assert_called_once_with(self.config)
+        self.assertEqual(vision_system._backend, mock_backend_instance)
 
-    @patch('platform.system')
+    @patch('platform.system', return_value='Darwin')
     @patch('vision.PythonBackend')
-    def test_initialization_macos(self, mock_backend, mock_system):
+    def test_initialization_macos(self, mock_backend_class, mock_system):
         """
         Test that the VisionSystem initializes the PythonBackend on macOS.
         """
-        mock_system.return_value = 'Darwin'
+        mock_backend_instance = mock_backend_class.return_value
         vision_system = VisionSystem(self.config)
-        mock_backend.assert_called_once_with(self.config)
+        mock_backend_class.assert_called_once_with(self.config)
+        self.assertEqual(vision_system._backend, mock_backend_instance)
 
-    @patch('platform.system')
+    @patch('platform.system', return_value='UnsupportedOS')
     def test_initialization_unsupported(self, mock_system):
         """
         Test that the VisionSystem raises an exception on an unsupported OS.
         """
-        mock_system.return_value = 'UnsupportedOS'
         with self.assertRaises(NotImplementedError):
             VisionSystem(self.config)
 
-    @patch('platform.system')
+    @patch('platform.system', return_value='Linux')
     @patch('vision.DeepStreamBackend')
-    def test_start(self, mock_backend, mock_system):
+    def test_start(self, mock_backend_class, mock_system):
         """
         Test that the start method calls the backend's start method.
         """
-        mock_system.return_value = 'Linux'
+        mock_backend_instance = mock_backend_class.return_value
         vision_system = VisionSystem(self.config)
         vision_system.start()
-        vision_system._backend.start.assert_called_once()
+        mock_backend_instance.start.assert_called_once()
 
-    @patch('platform.system')
+    @patch('platform.system', return_value='Linux')
     @patch('vision.DeepStreamBackend')
-    def test_stop(self, mock_backend, mock_system):
+    def test_stop(self, mock_backend_class, mock_system):
         """
         Test that the stop method calls the backend's stop method.
         """
-        mock_system.return_value = 'Linux'
+        mock_backend_instance = mock_backend_class.return_value
         vision_system = VisionSystem(self.config)
         vision_system.stop()
-        vision_system._backend.stop.assert_called_once()
+        mock_backend_instance.stop.assert_called_once()
 
-    @patch('platform.system')
+    @patch('platform.system', return_value='Linux')
     @patch('vision.DeepStreamBackend')
-    def test_get_pupil_data(self, mock_backend, mock_system):
+    def test_get_pupil_data(self, mock_backend_class, mock_system):
         """
         Test that the get_pupil_data method calls the backend's get_pupil_data method.
         """
-        mock_system.return_value = 'Linux'
+        mock_backend_instance = mock_backend_class.return_value
         vision_system = VisionSystem(self.config)
         vision_system.get_pupil_data()
-        vision_system._backend.get_pupil_data.assert_called_once()
+        mock_backend_instance.get_pupil_data.assert_called_once()
+
+    @patch('platform.system', return_value='Linux')
+    @patch('vision.DeepStreamBackend')
+    def test_get_annotated_frame(self, mock_backend_class, mock_system):
+        """
+        Test that the get_annotated_frame method calls the backend's get_annotated_frame method.
+        """
+        mock_backend_instance = mock_backend_class.return_value
+        vision_system = VisionSystem(self.config)
+        vision_system.get_annotated_frame()
+        mock_backend_instance.get_annotated_frame.assert_called_once()
+
+    @patch('vision.logging')
+    @patch.dict('sys.modules', {'onnxruntime': MagicMock(), 'onnxruntime-gpu': None})
+    def test_python_backend_cpu_fallback(self, mock_logging):
+        """
+        Test that PythonBackend falls back to CPU when onnxruntime-gpu is not available.
+        """
+        mock_ort = sys.modules['onnxruntime']
+        mock_ort.get_available_providers.return_value = ['CPUExecutionProvider']
+        
+        backend = PythonBackend(self.config)
+        
+        mock_logging.warning.assert_called_with("onnxruntime-gpu is not available or CUDA is not configured. Falling back to onnxruntime (CPU).")
+        self.assertEqual(backend.ort, mock_ort)
+
+    @patch('vision.logging')
+    @patch.dict('sys.modules', {'onnxruntime': MagicMock()})
+    def test_python_backend_gpu_selection(self, mock_logging):
+        """
+        Test that PythonBackend selects GPU when onnxruntime-gpu is available.
+        """
+        mock_ort_gpu = MagicMock()
+        mock_ort_gpu.get_available_providers.return_value = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        sys.modules['onnxruntime'] = mock_ort_gpu
+        
+        backend = PythonBackend(self.config)
+        
+        mock_logging.info.assert_any_call("CUDA is available. Using onnxruntime-gpu.")
+        self.assertEqual(backend.ort, mock_ort_gpu)
+
+    def test_mock_backend_methods(self):
+        """
+        Test the methods of the MockBackend.
+        """
+        backend = MockBackend(self.config)
+        backend.start()
+        backend.stop()
+        data = backend.get_pupil_data()
+        self.assertIn("pupil_position", data)
+        frame = backend.get_annotated_frame()
+        self.assertIsInstance(frame, np.ndarray)
 
 if __name__ == '__main__':
     unittest.main()