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.
This commit is contained in:
parent
8aebeea6ee
commit
40b9b2c8d2
@ -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"
|
||||
|
||||
8
run.ps1
Normal file
8
run.ps1
Normal file
@ -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
|
||||
4
run.sh
Normal file
4
run.sh
Normal file
@ -0,0 +1,4 @@
|
||||
#!/bin/bash
|
||||
source .venv/bin/activate
|
||||
pip install -r requirements.txt
|
||||
python 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
|
||||
# =================================================================================================
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
});
|
||||
@ -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);
|
||||
|
||||
@ -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>
|
||||
|
||||
@ -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
|
||||
print("Initializing DeepStream backend...")
|
||||
logging.info(f"Operating system is {os_name}. Attempting to initialize 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:
|
||||
print(f"Could not initialize DeepStreamBackend: {e}")
|
||||
raise e
|
||||
elif os_name == "Darwin":
|
||||
# On macOS, use the Python-based backend
|
||||
print("Initializing Python backend for macOS...")
|
||||
try:
|
||||
except (ImportError, ValueError) as e:
|
||||
logging.warning(f"Could not initialize DeepStreamBackend: {e}. Falling back to PythonBackend.")
|
||||
return PythonBackend(self.config)
|
||||
except ImportError as e:
|
||||
print(f"Could not initialize PythonBackend: {e}")
|
||||
raise e
|
||||
elif os_name == "Darwin":
|
||||
logging.info("Operating system is macOS. Initializing Python backend.")
|
||||
return PythonBackend(self.config)
|
||||
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'])
|
||||
print(f"PythonBackend: ONNX model loaded from {self.config['model_path']}.")
|
||||
self.inference_session = self.ort.InferenceSession(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.")
|
||||
print("PythonBackend stopped.")
|
||||
logging.info("PythonBackend: Basler camera closed.")
|
||||
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
|
||||
@ -254,18 +289,31 @@ 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}")
|
||||
|
||||
|
||||
@ -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:
|
||||
stdout = f.read()
|
||||
with open(STDERR_FILE, "r") as f:
|
||||
stderr = f.read()
|
||||
|
||||
assert "Initializing Mock backend for testing..." in stdout
|
||||
assert "MockBackend initialized." in stdout
|
||||
assert "Initializing Mock backend" in stderr
|
||||
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"
|
||||
|
||||
@ -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()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user