From a449d41a61160dbe4a94f02337e65c09d7aeb38f Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sat, 16 Aug 2025 02:46:19 +0000 Subject: [PATCH 1/2] feat: Implement major improvements based on repository audit This commit introduces a wide range of improvements based on the findings of a repository audit. The changes address key areas including testing, continuous integration, documentation, logging, performance, and extensibility. Key changes include: - **CI/CD:** Added a GitHub Actions workflow for automated linting and testing. - **Documentation:** Overhauled the README with setup instructions, usage guidance, and an architecture diagram. - **Structured Logging:** Refactored the entire codebase to use `structlog` for structured, JSON-formatted logging, replacing previous `logging` and `print` calls. - **Performance:** Enabled Automatic Mixed Precision (AMP) for YOLOv8 inference on CUDA devices to accelerate performance. - **Extensibility:** Implemented a plugin-based architecture for the object detector. The YOLOv8 and OpenVINO backends have been refactored into separate plugins, making the system more modular and extensible. - **Code Quality:** Added a comprehensive `.gitignore` file and fixed a large number of linting errors reported by `ruff`. The test suite was also updated to support the new architecture and is fully passing. --- .github/workflows/ci.yml | 44 +++ .gitignore | 121 +++++- README.md | 70 +++- poetry.lock | 14 +- pyproject.toml | 1 + src/zebtrack/__main__.py | 45 ++- src/zebtrack/core/controller.py | 109 +++--- src/zebtrack/core/detector.py | 346 +++--------------- src/zebtrack/core/project_manager.py | 108 +++--- src/zebtrack/io/arduino.py | 74 ++-- src/zebtrack/io/camera.py | 40 +- src/zebtrack/io/recorder.py | 103 ++---- src/zebtrack/io/sources.py | 3 +- src/zebtrack/io/video_source.py | 62 +--- src/zebtrack/plugins/__init__.py | 9 + src/zebtrack/plugins/base.py | 54 +++ src/zebtrack/plugins/openvino_detector.py | 133 +++++++ src/zebtrack/plugins/yolo_detector.py | 54 +++ src/zebtrack/settings.py | 47 ++- src/zebtrack/ui/gui.py | 122 +++--- src/zebtrack/utils.py | 6 +- tests/__pycache__/__init__.cpython-312.pyc | Bin 119 -> 119 bytes ...oject_manager.cpython-312-pytest-8.4.1.pyc | Bin 7728 -> 7709 bytes ...test_recorder.cpython-312-pytest-8.4.1.pyc | Bin 8893 -> 8851 bytes tests/test_controller.py | 1 + tests/test_detector.py | 163 +++------ tests/test_recorder.py | 1 - tests/test_settings.py | 2 - tests/test_sources.py | 2 +- 29 files changed, 969 insertions(+), 765 deletions(-) create mode 100644 .github/workflows/ci.yml create mode 100644 src/zebtrack/plugins/__init__.py create mode 100644 src/zebtrack/plugins/base.py create mode 100644 src/zebtrack/plugins/openvino_detector.py create mode 100644 src/zebtrack/plugins/yolo_detector.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..bb462d1 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,44 @@ +name: CI + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.12' + - name: Install Poetry + run: | + pip install poetry + - name: Install dependencies + run: | + poetry install --no-root + - name: Run ruff + run: | + poetry run ruff check . + + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.12' + - name: Install Poetry + run: | + pip install poetry + - name: Install dependencies + run: | + poetry install --no-root + - name: Run tests + run: | + poetry run pytest diff --git a/.gitignore b/.gitignore index 72701b6..bfb571c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,13 +1,122 @@ -# Python bytecode +# Byte-compiled / optimized / DLL files __pycache__/ -*.pyc +*.py[cod] +*$py.class -# PyTorch model files -*.pt +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST -# Test artifacts +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ .pytest_cache/ test_projects/ -# Log files +# Translations +*.mo +*.pot + +# Django stuff: *.log +local_settings.py +db.sqlite3 + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyderworkspace + +# Rope project settings +.ropeproject + +# PyCharm +.idea/ +.project +.pydevproject +.settings/ +*.iml +*.iws + +# VSCode +.vscode/ + +#mypy +.mypy_cache/ + +# PyTorch model files +*.pt +*.pth diff --git a/README.md b/README.md index eefe387..65fd537 100644 --- a/README.md +++ b/README.md @@ -4,11 +4,77 @@ Zebtrack Controller is a graphical application designed for object tracking in v ## Installation -_Instructions on how to install the necessary dependencies and set up the environment will be added here._ +This project is managed with [Poetry](https://python-poetry.org/). + +1. **Clone the repository:** + ```bash + git clone https://github.com/your-username/zebtrack-controller.git + cd zebtrack-controller + ``` + +2. **Install Poetry:** + Follow the official instructions at [python-poetry.org](https://python-poetry.org/docs/#installation) to install Poetry on your system. + +3. **Install dependencies:** + Once Poetry is installed, run the following command in the project root to create a virtual environment and install the required dependencies: + ```bash + poetry install + ``` ## Usage -_Instructions on how to run the application and use its features will be added here._ +To run the application, use the following command from the project's root directory: + +```bash +poetry run python -m zebtrack +``` + +This will launch the main graphical user interface. + +## Architecture + +The application is designed with a separation of concerns, loosely following a Model-View-Controller (MVC) pattern. + +```mermaid +graph TD + subgraph "User Interface (View)" + GUI["GUI (Tkinter)"] + end + + subgraph "Core Logic (Controller & Model)" + AppController + ProjectManager + Detector["Detector (Ultralytics/OpenVINO)"] + Settings + end + + subgraph "I/O Subsystem" + FrameSource["FrameSource (Camera/Video)"] + Recorder + Arduino + end + + GUI -- User Actions --> AppController + AppController -- Updates --> GUI + + AppController -- Manages --> ProjectManager + AppController -- Uses --> Settings + AppController -- Controls --> Detector + AppController -- Controls --> Recorder + AppController -- Controls --> Arduino + AppController -- Gets Frames --> FrameSource + + Detector -- Processes frames provided by --> AppController +``` + +* **GUI**: The user interface, built with Tkinter. +* **AppController**: The central component that handles user input from the GUI and coordinates all other components. +* **ProjectManager**: Manages the creation, loading, and saving of project files and configurations. +* **Detector**: Performs object detection on video frames using models from `ultralytics` or `OpenVINO`. +* **FrameSource**: Provides video frames, either from a live camera feed or a video file. +* **Recorder**: Handles the saving of output video and tracking data. +* **Arduino**: Manages communication with an Arduino board for hardware I/O. +* **Settings**: Loads and manages application settings from configuration files. ## License diff --git a/poetry.lock b/poetry.lock index 012d24a..255fe83 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1791,6 +1791,18 @@ files = [ {file = "six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81"}, ] +[[package]] +name = "structlog" +version = "25.4.0" +description = "Structured Logging for Python" +optional = false +python-versions = ">=3.8" +groups = ["main"] +files = [ + {file = "structlog-25.4.0-py3-none-any.whl", hash = "sha256:fe809ff5c27e557d14e613f45ca441aabda051d119ee5a0102aaba6ce40eed2c"}, + {file = "structlog-25.4.0.tar.gz", hash = "sha256:186cd1b0a8ae762e29417095664adf1d6a31702160a46dacb7796ea82f7409e4"}, +] + [[package]] name = "sympy" version = "1.14.0" @@ -2076,4 +2088,4 @@ zstd = ["zstandard (>=0.18.0)"] [metadata] lock-version = "2.1" python-versions = ">=3.12" -content-hash = "598c6aceaef3ac3296d7ef4ff8ff3042ef62cc67a7ff8e51cb9ff022815fbfd5" +content-hash = "a432de09ca1edf0b88ddbbdd49bda031b56a85bdef92ed9915930f3b8d56e122" diff --git a/pyproject.toml b/pyproject.toml index b544ef1..34606b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ openvino = "==2025.1.0" ultralytics = ">=8.3.179,<9.0.0" pyyaml = "^6.0.2" pydantic = "^2.11.7" +structlog = "^25.4.0" [tool.poetry.group.dev.dependencies] ruff = "^0.12.9" diff --git a/src/zebtrack/__main__.py b/src/zebtrack/__main__.py index 98bcbd2..a3ef400 100644 --- a/src/zebtrack/__main__.py +++ b/src/zebtrack/__main__.py @@ -1,6 +1,8 @@ import logging import tkinter as tk +import structlog + from zebtrack.core.controller import AppController from zebtrack.settings import settings from zebtrack.utils import set_seed @@ -10,30 +12,53 @@ def main(): """ Initializes and runs the application. """ - # Configure logging + # Configure logging with structlog + structlog.configure( + processors=[ + structlog.stdlib.add_log_level, + structlog.stdlib.add_logger_name, + structlog.processors.TimeStamper(fmt="iso", utc=True), + structlog.processors.StackInfoRenderer(), + structlog.processors.format_exc_info, + structlog.processors.UnicodeDecoder(), + structlog.processors.JSONRenderer() + ], + context_class=dict, + logger_factory=structlog.stdlib.LoggerFactory(), + wrapper_class=structlog.stdlib.BoundLogger, + cache_logger_on_first_use=True, + ) + + # Route standard library logs to structlog logging.basicConfig( level=logging.INFO, - format="%(asctime)s - %(threadName)s - %(name)s - %(levelname)s - %(message)s", - filename="analysis.log", - filemode="w", # Overwrite log file on each run + format="%(message)s", + handlers=[logging.FileHandler("analysis.log", mode="w")], ) + log = structlog.get_logger() + # Set seed for reproducibility before anything else - if settings and settings.reproducibility: + if settings and settings.reproducibility and settings.reproducibility.seed: set_seed(settings.reproducibility.seed) + log.info("reproducibility.seed.set", seed=settings.reproducibility.seed) - logging.info("Application starting.") + log.info("application.starting", component="main") try: root = tk.Tk() controller = AppController(root) controller.run() - except Exception as e: - logging.critical("An unhandled exception occurred.", exc_info=True) + except Exception: + log.critical("unhandled.exception", exc_info=True) # Optionally, show a message to the user - # messagebox.showerror("Fatal Error", f"A fatal error occurred: {e}\nSee analysis.log for details.") + # import tkinter.messagebox as messagebox + # messagebox.showerror( + # "Fatal Error", + # "A fatal error occurred. See analysis.log for details." + # ) finally: - logging.info("Application finished.") + log.info("application.finished", component="main") if __name__ == "__main__": diff --git a/src/zebtrack/core/controller.py b/src/zebtrack/core/controller.py index da763d3..4288cd1 100644 --- a/src/zebtrack/core/controller.py +++ b/src/zebtrack/core/controller.py @@ -6,12 +6,11 @@ and the backend modules (Model). """ -import logging import os import queue import threading -import time -from tkinter import messagebox + +import structlog from zebtrack.core.project_manager import ProjectManager from zebtrack.io.arduino import Arduino @@ -19,6 +18,8 @@ from zebtrack.io.recorder import Recorder from zebtrack.ui.gui import ApplicationGUI +log = structlog.get_logger() + class AppController: """ @@ -61,26 +62,23 @@ def __init__(self, root): self.processing_thread = None self.video_thread = None - logging.info("AppController initialized with backend modules and state.") + log.info("controller.init.success") def run(self): """ Starts the application's main loop. """ - logging.info("Starting main UI loop.") + log.info("ui.mainloop.start") self.root.mainloop() def on_close(self): """Handles the application shutdown process.""" - logging.info("Close button clicked.") + log.info("ui.close_button.clicked") if self.view.ask_ok_cancel("Quit", "Do you want to exit the program?"): - logging.info("User confirmed quit.") + log.info("user.quit.confirmed") self.program_exit_event.set() - # if self.is_recording and self.project_manager.get_project_type() == "live": - # self.stop_recording() # This will be moved later - self.join_threads() if self.camera: @@ -92,11 +90,11 @@ def on_close(self): self.arduino.close() self.root.destroy() - logging.info("Application shutdown complete.") + log.info("application.shutdown.complete") def join_threads(self): """Waits for all core threads to finish.""" - logging.info("Waiting for core threads to join.") + log.info("threads.join.start") if ( hasattr(self, "capture_thread") and self.capture_thread @@ -115,7 +113,7 @@ def join_threads(self): and self.video_thread.is_alive() ): self.video_thread.join(timeout=5) - logging.info("Core threads joined.") + log.info("threads.join.finished") def stop_recording(self): """Stops the recording for a 'live' project.""" @@ -123,21 +121,21 @@ def stop_recording(self): self.is_capturing_for_video = False self.video_stop_event.set() - # We need to join the video_thread here if hasattr(self, "video_thread") and self.video_thread.is_alive(): self.video_thread.join(timeout=5) self.recorder.stop_recording() - # Update the view self.view.update_button_state("start_rec", "normal") self.view.update_button_state("stop_rec", "disabled") - self.view.set_status(f"Project: {self.project_manager.get_project_name()} (live) - Ready") + self.view.set_status( + f"Project: {self.project_manager.get_project_name()} (live) - Ready" + ) self.view.show_info("Success", "Recording stopped and files saved.") def close_project(self): """Closes the current project and returns to the welcome screen.""" - logging.info("Closing project.") + log.info("project.close.start") self.program_exit_event.set() if self.is_recording and self.project_manager.get_project_type() == "live": @@ -153,14 +151,14 @@ def close_project(self): self.active_frame_source.release() self.active_frame_source = None - # Reset the project manager for the next project self.project_manager = ProjectManager() - # Tell the view to go back to the welcome screen self.view.create_welcome_frame() - logging.info("Project closed and welcome screen recreated.") + log.info("project.close.finished") - def create_project_workflow(self, project_path, project_type, use_openvino, video_files): + def create_project_workflow( + self, project_path, project_type, use_openvino, video_files + ): """Handles the logic of creating a new project.""" success = self.project_manager.create_new_project( project_path, @@ -169,8 +167,6 @@ def create_project_workflow(self, project_path, project_type, use_openvino, vide video_files=video_files, ) if success: - # This method still lives in the view, we need to move it. - # For now, we call it on the view. self.view._load_project_view() else: self.view.show_error("Error", "Failed to create the new project.") @@ -180,18 +176,24 @@ def open_project_workflow(self, project_path): if self.project_manager.load_project(project_path): self.view._load_project_view() else: - self.view.show_error("Error", "Failed to load the project. Check if it's a valid project folder.") + self.view.show_error( + "Error", + "Failed to load the project. Check if it's a valid project folder.", + ) def start_recording(self): """Handles the business logic for starting a recording session.""" if not self.project_manager.project_data.get("groups"): - self.view.show_warning("Setup Required", "Please define groups for this project first.") + self.view.show_warning( + "Setup Required", "Please define groups for this project first." + ) return - # Ask the view to get details from the user - details = self.view.ask_recording_details(self.project_manager.project_data["groups"]) + details = self.view.ask_recording_details( + self.project_manager.project_data["groups"] + ) if not details: - return # User cancelled + return group_name, cobaia_number = details output_folder = os.path.join( @@ -217,7 +219,6 @@ def start_recording(self): ) self.video_thread.start() - # Update the view self.view.update_button_state("start_rec", "disabled") self.view.update_button_state("stop_rec", "normal") self.view.set_status(f"Recording to: {os.path.basename(output_folder)}") @@ -228,14 +229,14 @@ def _video_recording_loop(self): """ Loop executed in a thread to write video frames to a file. """ - logging.info("Video recording thread started.") + log.info("video_thread.start") while not self.video_stop_event.is_set(): try: frame = self.video_queue.get(timeout=1) self.recorder.write_video_frame(frame) except queue.Empty: continue - logging.info("Video recording thread finished.") + log.info("video_thread.finished") def process_next_video(self): """Handles the business logic for processing the next pre-recorded video.""" @@ -245,14 +246,20 @@ def process_next_video(self): video_path = self.project_manager.get_next_video() if not video_path: - self.view.show_info("Project Complete", "All videos in this project have been processed.") + self.view.show_info( + "Project Complete", "All videos in this project have been processed." + ) return if not self.project_manager.project_data.get("groups"): - self.view.show_warning("Setup Required", "Please define groups for this project first.") + self.view.show_warning( + "Setup Required", "Please define groups for this project first." + ) return - details = self.view.ask_recording_details(self.project_manager.project_data["groups"]) + details = self.view.ask_recording_details( + self.project_manager.project_data["groups"] + ) if not details: return @@ -269,7 +276,9 @@ def process_next_video(self): video_basename = os.path.splitext(os.path.basename(video_path))[0] output_folder_name = f"{video_basename}_{group_name}_{cobaia_number}" - output_path = os.path.join(self.project_manager.project_path, output_folder_name) + output_path = os.path.join( + self.project_manager.project_path, output_folder_name + ) success = self.recorder.start_recording( output_path, video_props["width"], video_props["height"], is_video_file=True @@ -289,17 +298,20 @@ def process_next_video(self): self.view.update_button_state("process_video", "disabled") self.view.set_status(f"Processing: {os.path.basename(video_path)}") else: - self.view.show_error("Error", "Failed to start recorder for video processing.") + self.view.show_error( + "Error", "Failed to start recorder for video processing." + ) video_source.release() def _file_processing_loop(self): """ Loop for processing a video file. This is the core logic that runs in a thread. """ - from zebtrack.settings import settings - from zebtrack.core.detector import draw_overlay import cv2 + from zebtrack.core.detector import draw_overlay + from zebtrack.settings import settings + show_preview = self.view.show_preview_var.get() try: processing_interval = int(self.view.processing_interval_var.get()) @@ -313,11 +325,12 @@ def _file_processing_loop(self): frame_number = -1 while not self.program_exit_event.is_set() and frame_number < total_frames: - target_frame = ( - (settings.video_processing.processing_offset if settings.video_processing.processing_offset > 0 else 1) - if frame_number < 0 - else frame_number + processing_interval - ) + if frame_number < 0: + offset = settings.video_processing.processing_offset + target_frame = offset if offset > 0 else 1 + else: + target_frame = frame_number + processing_interval + if target_frame >= total_frames: break video_source.cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame) @@ -347,7 +360,7 @@ def _file_processing_loop(self): def cleanup_after_processing(self): """Cleans up resources after video processing is complete.""" video_name = os.path.basename(self.currently_processing_video) - logging.info(f"Cleaning up after processing video: {video_name}.") + log.info("video_processing.cleanup.start", video_name=video_name) self.is_recording = False self.recorder.stop_recording() @@ -365,7 +378,11 @@ def cleanup_after_processing(self): next_video = self.project_manager.get_next_video() if next_video: status_msg = f"Ready to process: {os.path.basename(next_video)}" - self.view.set_status(f"Project: {self.project_manager.get_project_name()} - {status_msg}") + self.view.set_status( + f"Project: {self.project_manager.get_project_name()} - {status_msg}" + ) else: status_msg = "All videos processed." - self.view.set_status(f"Project: {self.project_manager.get_project_name()} - {status_msg}") + self.view.set_status( + f"Project: {self.project_manager.get_project_name()} - {status_msg}" + ) diff --git a/src/zebtrack/core/detector.py b/src/zebtrack/core/detector.py index bc83bf7..69b30fd 100644 --- a/src/zebtrack/core/detector.py +++ b/src/zebtrack/core/detector.py @@ -1,87 +1,42 @@ -""" -Este módulo contém a classe Detector, responsável por usar o modelo YOLO -para detectar objetos em quadros de vídeo e a lógica associada para rastrear -a entrada e saída de áreas de interesse. -""" - -import glob -import logging -import os import time import cv2 import numpy as np -import openvino as ov -import torch -from ultralytics import YOLO -from ultralytics.utils.ops import non_max_suppression, scale_boxes +import structlog +from zebtrack.plugins.base import DetectorPlugin from zebtrack.settings import settings +log = structlog.get_logger() + class Detector: """ - Encapsula o modelo de detecção (YOLO ou OpenVINO) e a lógica de detecção. - - Esta classe carrega o modelo de detecção de objetos, gerencia as coordenadas - das áreas de interesse (escalando-as para a resolução do vídeo) e processa - quadros individuais para encontrar objetos. Para projetos 'live', também - implementa uma máquina de estados simples para gerar comandos para o Arduino - baseado na movimentação do objeto detectado entre as áreas. + Manages the detection process by delegating to a plugin and handling + stateful logic for zone tracking. """ - def __init__(self, project_manager=None, model_path: str = None): + def __init__(self, plugin: DetectorPlugin): """ - Inicializa o detector de objetos. - - - Carrega o modelo YOLO ou OpenVINO com base na configuração do projeto. - - Define os limiares de confiança e Non-Maximum Suppression (NMS). - - Inicializa variáveis de estado para rastrear a movimentação do objeto. - - Define as coordenadas das áreas de interesse com base na resolução padrão. + Initializes the detector with a specific plugin. Args: - project_manager: O gerenciador de projetos para configurações. - model_path (str, optional): Caminho para o modelo a ser carregado, - substituindo o padrão dos settings. + plugin (DetectorPlugin): An instantiated detector plugin. """ - self.model = None - self.is_openvino = False - self.compiled_model = None - self.input_layer = None - self.output_layer = None - self.infer_request = None + self.plugin = plugin + if not self.plugin: + log.error("detector.init.no_plugin") + raise ValueError("Detector must be initialized with a valid plugin.") - use_openvino = False - openvino_path = "" - if project_manager and project_manager.project_data: - use_openvino = project_manager.project_data.get("use_openvino", False) - openvino_path = project_manager.project_data.get("openvino_model_path", "") + log.info("detector.init.success", plugin=self.plugin.get_name()) - try: - if use_openvino and openvino_path: - self._load_openvino_model(openvino_path) - self.is_openvino = True - else: - # Usa o model_path fornecido ou o caminho dos settings - path_to_load = model_path or settings.yolo_model.path - logging.info(f"Loading YOLO model from: {path_to_load}") - self.model = YOLO(path_to_load) - except Exception as e: - logging.critical(f"Failed to load detection model: {e}") - # As variáveis de modelo permanecem None, desativando o detector - - self.conf_threshold = settings.yolo_model.confidence_threshold - self.nms_threshold = settings.yolo_model.nms_threshold - - # Variáveis de estado para rastrear a movimentação do objeto - # (usado em projetos 'live'). + # State variables for tracking object movement self.crossed_in = False self.crossed_out = False - self.flag = 0 # 0: procurando entrada, 1: procurando saída - self.current_square = 0 # Qual quadrado o objeto entrou + self.flag = 0 # 0: looking for entry, 1: looking for exit + self.current_square = 0 - # As coordenadas são definidas em `config.py` para uma resolução base - # e escaladas para a resolução real do vídeo pela `update_scaling`. + # Zone coordinates are defined in settings and scaled for the video resolution. self.base_polygon = np.array( settings.detection_zones.polygon, dtype=np.int32 ) @@ -92,22 +47,14 @@ def __init__(self, project_manager=None, model_path: str = None): settings.camera.desired_width, settings.camera.desired_height ) - def update_scaling(self, actual_width, actual_height): + def update_scaling(self, actual_width: int, actual_height: int): """ - Atualiza as coordenadas do polígono e dos quadrados com base na - resolução real do vídeo. - - Isso permite que as áreas de interesse sejam definidas uma vez em `config.py` - e funcionem com vídeos de diferentes tamanhos. - - Args: - actual_width (int): A largura real da fonte de vídeo. - actual_height (int): A altura real da fonte de vídeo. + Updates the coordinates of the polygon and squares based on the actual + video resolution. """ base_width = settings.camera.desired_width base_height = settings.camera.desired_height - # Se a resolução já for a base, não há necessidade de escalar. if actual_width == base_width and actual_height == base_height: self.scaled_polygon = self.base_polygon self.scaled_squares = self.base_squares @@ -116,12 +63,9 @@ def update_scaling(self, actual_width, actual_height): scale_x = actual_width / base_width scale_y = actual_height / base_height - # Escala o polígono self.scaled_polygon = (self.base_polygon * [scale_x, scale_y]).astype( np.int32 ) - - # Escala os quadrados self.scaled_squares = [] for p1, p2 in self.base_squares: x1, y1 = p1 @@ -130,149 +74,37 @@ def update_scaling(self, actual_width, actual_height): scaled_p2 = (int(x2 * scale_x), int(y2 * scale_y)) self.scaled_squares.append((scaled_p1, scaled_p2)) - logging.info( - f"Detector coordinates scaled for resolution " - f"{actual_width}x{actual_height}" + log.info( + "detector.scaling.updated", + width=actual_width, + height=actual_height, ) def _is_inside_square(self, x1, y1, x2, y2, square): - """Verifica se uma caixa delimitadora se sobrepõe a um quadrado de área.""" + """Checks if a bounding box overlaps with an area square.""" (sx1, sy1), (sx2, sy2) = square return not (x2 < sx1 or x1 > sx2 or y2 < sy1 or y1 > sy2) def _is_inside_polygon(self, x1, y1, x2, y2, polygon): - """Verifica se um canto da caixa delimitadora está dentro do polígono.""" + """Checks if a corner of the bounding box is inside the polygon.""" return ( cv2.pointPolygonTest(polygon, (x1, y1), False) >= 0 or cv2.pointPolygonTest(polygon, (x2, y2), False) >= 0 ) - def _load_openvino_model(self, model_dir_path): - """ - Loads the OpenVINO model from the specified directory. - It finds the .xml file within the directory to load the model. - """ - xml_files = glob.glob(os.path.join(model_dir_path, "*.xml")) - if not xml_files: - raise FileNotFoundError( - f"Could not find a .xml model file in directory: {model_dir_path}" - ) - - model_xml_path = xml_files[0] - logging.info(f"Found OpenVINO model file: {model_xml_path}") - - core = ov.Core() - model = core.read_model(model_xml_path) - self.compiled_model = core.compile_model(model=model, device_name="AUTO") - self.input_layer = self.compiled_model.input(0) - self.output_layer = self.compiled_model.output(0) - self.infer_request = self.compiled_model.create_infer_request() - - def _preprocess_openvino(self, frame): - """ - Prepares a frame for OpenVINO inference using letterboxing, which is - the standard for YOLO models. - """ - # Get input size of the model - n, c, h, w = self.input_layer.shape - - # Apply letterboxing. `auto=False` ensures the frame is padded - # to the exact `new_shape` (e.g., 640x640), required for static shapes. - letterboxed_frame, _, _ = letterbox(frame, new_shape=(w, h), auto=False) - - # Convert from BGR to RGB - rgb_frame = cv2.cvtColor(letterboxed_frame, cv2.COLOR_BGR2RGB) - - # Transpose from HWC to CHW and normalize to [0,1] - input_tensor = rgb_frame.transpose(2, 0, 1) / 255.0 - - # Add batch dimension to create NHWC - input_tensor = np.expand_dims(input_tensor, axis=0).astype(np.float32) - - return input_tensor - - def _postprocess_openvino(self, result, original_frame_shape): - """ - Postprocesses the OpenVINO model's output using the official - ultralytics utility functions for robust and accurate results. - """ - # Get the raw output tensor from the model - output_tensor = result[self.output_layer] - - # Use the official ultralytics non_max_suppression utility. - # This handles all the complex parsing of class scores and confidences. - # The output shape of the model is (1, 84, 8400) for COCO. - # We pass it directly to the utility. - preds = non_max_suppression( - prediction=torch.from_numpy(output_tensor), - conf_thres=self.conf_threshold, - iou_thres=self.nms_threshold, - agnostic=True, # Class-agnostic NMS - ) - - # The result of NMS is a list with one element per image in the batch. - # We only have one image, so we take the first element. - detections = preds[0] - - if detections is None or len(detections) == 0: - return [] - - # Rescale the bounding boxes from the model's input size (e.g., 640x640) - # back to the original frame's size. - model_input_shape = ( - self.input_layer.shape[2], - self.input_layer.shape[3], - ) # (h, w) - detections[:, :4] = scale_boxes( - model_input_shape, detections[:, :4], original_frame_shape - ).round() - - # Convert the results to the format expected by the rest of the application: - # A list of tuples: (x1, y1, x2, y2, confidence) - final_detections = [] - for *xyxy, conf, cls in detections: - # We ignore 'cls' for now as the application is class-agnostic - final_detections.append( - (int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3]), float(conf)) - ) - - return final_detections - - def process_frame(self, frame, project_type): + def process_frame(self, frame: np.ndarray, project_type: str): """ - Processa um único quadro para detecção de objetos e rastreamento de estado. + Processes a single frame for object detection and state tracking. """ - # Se o modelo não foi carregado com sucesso, não faz nada. - if self.model is None and self.compiled_model is None: - return [], None - start_time = time.perf_counter() - if self.is_openvino: - # --- OpenVINO Inference Path --- - input_tensor = self._preprocess_openvino(frame) - # The `infer` method is the recommended sync approach in the latest API - self.infer_request.infer({self.input_layer.any_name: input_tensor}) - results = self.infer_request.results - predictions = self._postprocess_openvino(results, frame.shape) - # The output is already in (x1, y1, x2, y2, confidence) format - else: - # --- Default Ultralytics Inference Path --- - results = self.model( - frame, verbose=False, conf=self.conf_threshold, iou=self.nms_threshold - ) - # Convert to a common format: list of (x1, y1, x2, y2, confidence) - predictions = [] - for det in results[0].boxes.data.cpu().numpy(): - x1, y1, x2, y2, confidence, _ = det - predictions.append((int(x1), int(y1), int(x2), int(y2), confidence)) + # 1. Delegate actual detection to the loaded plugin + predictions = self.plugin.detect(frame) - # --- Common Logic for both paths --- + # 2. Apply stateful logic based on the detections detections_in_polygon = [] command_to_send = None - found_object_for_state_change = ( - False # Garante que apenas um comando seja enviado por quadro - ) + found_object_for_state_change = False if len(predictions) > 0: for det in predictions: @@ -283,7 +115,7 @@ def process_frame(self, frame, project_type): detections_in_polygon.append((x1, y1, x2, y2, confidence)) if project_type == "live" and not found_object_for_state_change: - if self.flag == 0: + if self.flag == 0: # Looking for entry for index, square in enumerate(self.scaled_squares): if self._is_inside_square(x1, y1, x2, y2, square): self.crossed_in = True @@ -294,7 +126,7 @@ def process_frame(self, frame, project_type): ) found_object_for_state_change = True break - elif self.flag == 1: + elif self.flag == 1: # Looking for exit is_in_any_square = any( self._is_inside_square(x1, y1, x2, y2, sq) for sq in self.scaled_squares @@ -302,82 +134,33 @@ def process_frame(self, frame, project_type): if not is_in_any_square: self.crossed_out = True self.flag = 0 - command_to_send = settings.detection_zones.exit_commands[ - self.current_square - 1 - ] + command_to_send = ( + settings.detection_zones.exit_commands[ + self.current_square - 1 + ] + ) self.current_square = 0 found_object_for_state_change = True end_time = time.perf_counter() - logging.debug(f"Frame processing time: {(end_time - start_time) * 1000:.2f} ms") + log.debug( + "frame.processing.time", + duration_ms=(end_time - start_time) * 1000, + plugin=self.plugin.get_name(), + ) return detections_in_polygon, command_to_send -def letterbox( - img: np.ndarray, - new_shape: tuple = (640, 640), - color: tuple = (114, 114, 114), - auto: bool = True, - scaleFill: bool = False, - scaleup: bool = True, - stride: int = 32, -): - """ - Resize and pad image while meeting stride-multiple constraints. - This is the standard letterboxing function from the ultralytics library. - """ - shape = img.shape[:2] # current shape [height, width] - if isinstance(new_shape, int): - new_shape = (new_shape, new_shape) - - # Scale ratio (new / old) - r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) - if not scaleup: # only scale down, do not scale up (for better test mAP) - r = min(r, 1.0) - - # Compute padding - ratio = r, r # width, height ratios - new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) - dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding - if auto: # minimum rectangle - dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding - elif scaleFill: # stretch - dw, dh = 0.0, 0.0 - new_unpad = (new_shape[1], new_shape[0]) - ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios - - dw /= 2 # divide padding into 2 sides - dh /= 2 - - if shape[::-1] != new_unpad: # resize - img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) - top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) - left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) - img = cv2.copyMakeBorder( - img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color - ) # add border - return img, ratio, (dw, dh) - - def draw_overlay(frame, detections, detector_instance): """ - Desenha sobreposições de detecção no quadro. - - Isso inclui as áreas de interesse, o polígono de processamento e as - caixas delimitadoras para cada objeto detectado. - - Args: - frame: O quadro de vídeo no qual desenhar. - detections (list): A lista de detecções a serem desenhadas. - detector_instance (Detector): A instância do detector que contém as - coordenadas escaladas das áreas. + Draws detection overlays on the frame. """ - # Desenha os quadrados de área de interesse + # Draw the area-of-interest squares for i, ((x1, y1), (x2, y2)) in enumerate(detector_instance.scaled_squares): cv2.rectangle(frame, (x1, y1), (x2, y2), settings.detection_zones.colors[i], 2) - # Desenha o polígono da área de processamento + # Draw the processing area polygon cv2.polylines( frame, [detector_instance.scaled_polygon], @@ -386,7 +169,7 @@ def draw_overlay(frame, detections, detector_instance): thickness=1, ) - # Desenha as caixas delimitadoras das detecções + # Draw the bounding boxes for detections for x1, y1, x2, y2, confidence in detections: cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 255), 2) cv2.putText( @@ -398,34 +181,3 @@ def draw_overlay(frame, detections, detector_instance): (255, 0, 255), 2, ) - - -if __name__ == "__main__": - # This test requires a camera and will display the output. - from zebtrack.io.camera import Camera - - print("Running detector test...") - cam = Camera() - detector = Detector() - - while True: - ret, frame = cam.get_frame() - if not ret: - print("Failed to get frame.") - break - - detections, command = detector.process_frame(frame, "live") - - if command is not None: - print(f"Detector generated command: {command}") - - draw_overlay(frame, detections, detector) - - cv2.imshow("Detector Test", frame) - - if cv2.waitKey(1) & 0xFF == ord("q"): - break - - cam.release() - cv2.destroyAllWindows() - print("Detector test finished.") diff --git a/src/zebtrack/core/project_manager.py b/src/zebtrack/core/project_manager.py index d1f3575..800f567 100644 --- a/src/zebtrack/core/project_manager.py +++ b/src/zebtrack/core/project_manager.py @@ -1,10 +1,10 @@ import hashlib import json -import logging import os import shutil from tkinter import messagebox +import structlog import yaml from ultralytics import YOLO @@ -13,18 +13,19 @@ CONFIG_FILE_NAME = "project_config.json" SETTINGS_SNAPSHOT_FILE_NAME = "config_snapshot.yaml" +log = structlog.get_logger() + def _calculate_sha256(filepath: str) -> str: """Calculates the SHA256 hash of a file.""" sha256_hash = hashlib.sha256() try: with open(filepath, "rb") as f: - # Read and update hash in chunks of 4K for byte_block in iter(lambda: f.read(4096), b""): sha256_hash.update(byte_block) return sha256_hash.hexdigest() except IOError: - logging.error(f"Could not read file to calculate hash: {filepath}") + log.error("file.hash.read_error", filepath=filepath) return "" @@ -40,15 +41,13 @@ def _save_settings_snapshot(self): snapshot_path = os.path.join(self.project_path, SETTINGS_SNAPSHOT_FILE_NAME) try: - # Convert Pydantic settings to a dict and save as YAML settings_dict = settings.model_dump(mode="json") with open(snapshot_path, "w") as f: yaml.dump(settings_dict, f, indent=4, sort_keys=False) - logging.info(f"Settings snapshot saved to {snapshot_path}") + log.info("settings.snapshot.saved", path=snapshot_path) return True except (IOError, TypeError) as e: - # Log the error but don't block project creation - logging.error(f"Could not save settings snapshot: {e}") + log.error("settings.snapshot.save_error", error=str(e)) return False def create_new_project( @@ -59,6 +58,12 @@ def create_new_project( If use_openvino is True, it also converts the model to OpenVINO format. """ self.project_path = project_path + log_context = log.bind( + project_path=project_path, + project_type=project_type, + use_openvino=use_openvino, + ) + log_context.info("project.create.start") if project_type == "pre-recorded" and not video_files: raise ValueError("Pre-recorded projects require a list of video files.") @@ -66,75 +71,70 @@ def create_new_project( try: os.makedirs(self.project_path, exist_ok=True) except OSError as e: + log.error("project.create.dir_error", error=str(e)) messagebox.showerror( "Creation Error", - f"Could not create project directory:\n{e}\n\n" - "Please check folder permissions and ensure the path is valid.", + ( + f"Could not create project directory:\n{e}\n\n" + "Please check folder permissions and ensure the path is valid." + ), ) return False - # Save a snapshot of the settings for reproducibility self._save_settings_snapshot() - # Export the model if requested, checking for a cached version first. openvino_model_path = "" if use_openvino: cache_dir = "openvino_model_cache" base_model_name = os.path.splitext( os.path.basename(settings.yolo_model.path) )[0] - # Ultralytics appends `_openvino_model` to the exported directory name. cached_model_dir_name = f"{base_model_name}_openvino_model" cached_model_dir = os.path.join(cache_dir, cached_model_dir_name) - # The actual model file has .xml extension - # Note: ultralytics export might name it `best.xml` or - # `.xml`. We will just check for existence of the - # directory for simplicity, detector.py finds the xml. - if os.path.exists(cached_model_dir): - logging.info(f"Found cached OpenVINO model at {cached_model_dir}") + log.info( + "openvino.cache.found", + path=cached_model_dir, + ) openvino_model_path = os.path.abspath(cached_model_dir) else: - logging.info("No cached OpenVINO model found. Exporting now...") + log.info("openvino.export.start") try: model = YOLO(settings.yolo_model.path) - # Export to a temporary default location exported_path = model.export(format="openvino", half=True) - - # Ensure cache directory exists os.makedirs(cache_dir, exist_ok=True) - - # Move the exported model to our cache directory. - # The name of the exported dir is returned by `exported_path`. - # To prevent race conditions, we remove the destination - # first if it exists. shutil.rmtree(cached_model_dir, ignore_errors=True) try: shutil.move(exported_path, cached_model_dir) openvino_model_path = os.path.abspath(cached_model_dir) - logging.info( - f"Model exported and cached at {openvino_model_path}" + log.info( + "openvino.export.success", + path=openvino_model_path, ) except Exception as move_exc: - logging.error( - "Failed to move exported model to cache directory: " - f"{move_exc}" + log.error( + "openvino.export.move_error", + exc_info=move_exc, ) messagebox.showerror( "OpenVINO Export Error", - "Failed to move the exported OpenVINO model to the cache directory.\n" - f"Please check permissions.\n\nError: {move_exc}", + ( + "Failed to move the exported OpenVINO model to the " + f"cache directory.\nPlease check permissions.\n\nError: {move_exc}" + ), ) return False except Exception as e: - logging.error(f"Failed to export model to OpenVINO format: {e}") + log.error("openvino.export.failed", exc_info=e) messagebox.showerror( "OpenVINO Export Error", - "An unexpected error occurred during OpenVINO model export.\n" - "Ensure all dependencies are installed correctly and the model path is valid.\n\n" - f"Error: {e}", + ( + "An unexpected error occurred during OpenVINO model export.\n" + "Ensure all dependencies are installed correctly and the model " + f"path is valid.\n\nError: {e}" + ), ) return False @@ -153,7 +153,7 @@ def create_new_project( { "path": video_path, "sha256": video_hash, - "status": "pending", # Other statuses: "processing", "complete" + "status": "pending", } ) @@ -164,11 +164,18 @@ def load_project(self, project_path): Loads project data from a config file in the given directory. """ config_path = os.path.join(project_path, CONFIG_FILE_NAME) + log_context = log.bind(path=config_path) + log_context.info("project.load.start") + if not os.path.exists(config_path): + log_context.error("project.load.not_found") messagebox.showerror( "Load Error", - f"Project config file '{CONFIG_FILE_NAME}' not found in the selected directory:\n{project_path}\n\n" - "Please ensure you have selected a valid project folder.", + ( + f"Project config file '{CONFIG_FILE_NAME}' not found in the " + f"selected directory:\n{project_path}\n\nPlease ensure you have " + "selected a valid project folder." + ), ) return False @@ -176,12 +183,13 @@ def load_project(self, project_path): with open(config_path, "r") as f: self.project_data = json.load(f) self.project_path = project_path - print( - f"Project '{self.project_data.get('project_name')}' " - "loaded successfully." + log_context.info( + "project.load.success", + project_name=self.project_data.get("project_name"), ) return True except (json.JSONDecodeError, IOError) as e: + log_context.error("project.load.error", exc_info=e) messagebox.showerror( "Load Error", f"Failed to load or parse the project config file:\n{config_path}\n\n" @@ -200,9 +208,10 @@ def save_project(self): try: with open(config_path, "w") as f: json.dump(self.project_data, f, indent=4) - print(f"Project state saved to {config_path}") + log.info("project.save.success", path=config_path) return True except IOError as e: + log.error("project.save.error", path=config_path, exc_info=e) messagebox.showerror( "Save Error", f"Failed to save project config file:\n{config_path}\n\n" @@ -217,9 +226,10 @@ def update_video_status(self, video_path, new_status): for video in self.project_data.get("videos", []): if video["path"] == video_path: video["status"] = new_status - print( - f"Updated status of '{os.path.basename(video_path)}' to " - f"'{new_status}'" + log.info( + "video.status.update", + video_path=video_path, + status=new_status, ) return self.save_project() return False @@ -231,7 +241,7 @@ def get_next_video(self): for video in self.project_data.get("videos", []): if video["status"] == "pending": return video["path"] - return None # No more pending videos + return None def get_project_name(self): return self.project_data.get("project_name", "N/A") diff --git a/src/zebtrack/io/arduino.py b/src/zebtrack/io/arduino.py index 8b41045..0b5daba 100644 --- a/src/zebtrack/io/arduino.py +++ b/src/zebtrack/io/arduino.py @@ -1,65 +1,55 @@ -import logging import time from types import TracebackType from typing import Optional, Type import serial +import structlog from zebtrack.settings import settings +log = structlog.get_logger() + class Arduino: """ Manages serial communication with an Arduino device. - - This class handles connecting, sending commands, and receiving acknowledgments. - It is designed to be used as a context manager to ensure that the serial - connection is always closed properly. - - The expected communication protocol is as follows: - 1. On connection, the Arduino sends a ready message: "Arduino is ready.\\n". - 2. The host sends a command as an integer followed by a newline (e.g., "1\\n"). - 3. The Arduino processes the command and responds with "OK\\n" to acknowledge. """ def __init__(self, port: str, baud_rate: int): """ Initializes the Arduino controller. - Args: - port (str): The serial port the Arduino is connected to. - baud_rate (int): The baud rate for the serial communication. """ self.port = port self.baud_rate = baud_rate self.ser: Optional[serial.Serial] = None - logging.info("Arduino module initialized in offline mode.") + log.info("arduino.init", port=self.port, baud_rate=self.baud_rate) def connect(self) -> bool: """ - Attempts to establish a serial connection with the Arduino. It waits for a - "ready" signal from the Arduino after opening the serial port. - Returns True on success, False on failure. + Attempts to establish a serial connection with the Arduino. """ if self.ser and self.ser.is_open: - logging.info("Already connected to Arduino.") + log.info("arduino.connect.already_connected") return True try: self.ser = serial.Serial(self.port, self.baud_rate, timeout=2) - # The Arduino is expected to send a "ready" signal upon startup. - # This is more robust than a fixed sleep(). ready_signal = self.ser.readline().decode("utf-8").strip() if ready_signal == "Arduino is ready.": - logging.info(f"Successfully connected to Arduino on port {self.port}") + log.info("arduino.connect.success", port=self.port) return True else: - logging.warning(f"Arduino on port {self.port} did not send ready signal. " - f"Received: '{ready_signal}'") + log.warning( + "arduino.connect.no_ready_signal", + port=self.port, + received=ready_signal, + ) self.ser.close() self.ser = None return False except (serial.SerialException, OSError) as e: - logging.warning(f"Could not connect to Arduino on port {self.port}. {e}") - logging.warning("Running in offline mode. No commands will be sent to Arduino.") + log.warning( + "arduino.connect.failed", port=self.port, exc_info=e + ) self.ser = None return False @@ -81,66 +71,62 @@ def __exit__( def send_command(self, box_number: int) -> bool: """ Sends a command to the Arduino and waits for an acknowledgment. - Args: - box_number (int): The command number to send. - Returns: - bool: True if the command was sent and acknowledged, False otherwise. """ try: - # Ensure box_number is an integer command_num = int(box_number) except (ValueError, TypeError): - logging.error(f"Invalid command: '{box_number}' is not a valid integer.") + log.error("arduino.command.invalid", command=box_number) return False if self.ser and self.ser.is_open: command = f"{command_num}\n" try: self.ser.write(command.encode("utf-8")) - logging.info(f"Sent command to Arduino: {command.strip()}") + log.info("arduino.command.sent", command=command_num) - # Wait for acknowledgment response = self.ser.readline().decode("utf-8").strip() if response == "OK": - logging.info("Arduino acknowledged command.") + log.info("arduino.command.ack", command=command_num) return True else: - logging.warning( - f"Arduino acknowledgment failed. Expected 'OK', got '{response}'" + log.warning( + "arduino.command.nack", + command=command_num, + response=response, ) return False except serial.SerialException as e: - logging.error(f"Error during serial communication: {e}") + log.error("arduino.command.send_error", exc_info=e) return False else: - logging.debug(f"Offline mode: Command '{command_num}' not sent.") + log.debug("arduino.command.offline", command=command_num) return False def close(self) -> None: """ - Closes the serial connection and clears the serial object. + Closes the serial connection. """ if self.ser and self.ser.is_open: self.ser.close() - logging.info("Arduino connection closed.") + log.info("arduino.connection.closed") self.ser = None def main(): """Main function to run a test of the Arduino module.""" - logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") + # This is a test function, using print is fine here. print("Testing Arduino communication...") if not settings: print("Settings could not be loaded. Aborting test.") return - # Initialize and use Arduino with a context manager try: - with Arduino(port=settings.arduino.port, baud_rate=settings.arduino.baud_rate) as arduino: + with Arduino( + port=settings.arduino.port, baud_rate=settings.arduino.baud_rate + ) as arduino: print(f"Successfully connected to Arduino on {arduino.port}.") - # Test sending a sequence of commands print("\nSending test commands (1 to 8)...") all_commands_successful = True for i in range(1, 9): diff --git a/src/zebtrack/io/camera.py b/src/zebtrack/io/camera.py index 15e8bdf..1e9b8cc 100644 --- a/src/zebtrack/io/camera.py +++ b/src/zebtrack/io/camera.py @@ -1,19 +1,22 @@ +import threading +import time from typing import Any, Dict, Tuple + import cv2 import numpy as np -import threading -import time +import structlog from zebtrack.io.frame_source import FrameSource from zebtrack.settings import settings +log = structlog.get_logger() + class Camera(FrameSource): def __init__(self): self._camera_index = settings.camera.index self.cap = cv2.VideoCapture(self._camera_index) if not self.cap.isOpened(): - # This is a hard failure on startup, so we raise an exception raise IOError(f"Cannot open camera at index {self._camera_index}") self._desired_width = settings.camera.desired_width @@ -23,12 +26,13 @@ def __init__(self): self.actual_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) self.actual_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - print( - f"Camera initialized with resolution: " - f"{self.actual_width}x{self.actual_height}" + log.info( + "camera.initialized", + index=self._camera_index, + width=self.actual_width, + height=self.actual_height, ) - # Threading attributes self._lock = threading.Lock() self._latest_frame: Tuple[bool, np.ndarray | None] = (False, None) self._stopped = threading.Event() @@ -42,39 +46,36 @@ def _reader_thread(self): """ while not self._stopped.is_set(): if not self.cap.isOpened(): - print("Camera connection lost. Attempting to reconnect...") + log.warning("camera.reconnect.start") self.cap.open(self._camera_index) if self.cap.isOpened(): - print("Camera reconnected successfully.") + log.info("camera.reconnect.success") self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self._desired_width) self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self._desired_height) else: with self._lock: self._latest_frame = (False, None) - time.sleep(2) # Wait before next attempt + time.sleep(2) continue ret, frame = self.cap.read() if not ret: - # Frame read failed, likely due to a disconnect self.cap.release() - print("Frame read failed. Assuming disconnect, will try to reconnect.") + log.warning("camera.frame_read.failed") with self._lock: self._latest_frame = (False, None) - continue # Let the next loop iteration handle reconnection + continue with self._lock: self._latest_frame = (ret, frame) - print("Reader thread stopped.") + log.info("camera.reader_thread.stopped") def get_frame(self) -> Tuple[bool, np.ndarray | None]: """ Returns the most recent frame read by the background thread. - This method is non-blocking. """ with self._lock: - # Return a copy to prevent race conditions if the consumer modifies the frame ret, frame = self._latest_frame return ret, frame.copy() if ret else None @@ -83,10 +84,10 @@ def release(self) -> None: Signals the reader thread to stop and releases the camera resource. """ self._stopped.set() - self._thread.join(timeout=2) # Wait for the thread to finish + self._thread.join(timeout=2) if self.cap.isOpened(): self.cap.release() - print("Camera released.") + log.info("camera.released") def get_properties(self) -> Dict[str, Any]: """ @@ -106,14 +107,13 @@ def get_properties(self) -> Dict[str, Any]: camera = Camera() print("Camera properties:", camera.get_properties()) - # Give the reader thread a moment to start and grab the first frame time.sleep(1) while True: ret, frame = camera.get_frame() if not ret: print("Failed to grab frame, waiting...") - time.sleep(0.5) # Wait a bit before trying again + time.sleep(0.5) continue cv2.imshow("Camera Test", frame) diff --git a/src/zebtrack/io/recorder.py b/src/zebtrack/io/recorder.py index 70f2832..f6953e9 100644 --- a/src/zebtrack/io/recorder.py +++ b/src/zebtrack/io/recorder.py @@ -1,25 +1,23 @@ import csv -import logging import os import time import cv2 import numpy as np +import structlog from zebtrack.settings import settings +log = structlog.get_logger() + class Recorder: """ - Gerencia a gravação dos dados da análise, incluindo arquivos de vídeo e CSV. - - Esta classe lida com a criação de arquivos de saída, escrita de dados de - detecção e de quadros de vídeo, e o fechamento adequado dos arquivos ao - final da gravação. + Manages the recording of analysis data, including video and CSV files. """ def __init__(self): - """Inicializa o gravador com seu estado padrão.""" + """Initializes the recorder with its default state.""" self.is_recording = False self.video_writer = None self.csv_writer = None @@ -33,29 +31,27 @@ def start_recording( self, output_folder, frame_width, frame_height, is_video_file=False ): """ - Prepara e inicia uma nova sessão de gravação. - - Cria o diretório de saída e inicializa os escritores de CSV e, - opcionalmente, de vídeo. Salva também os arquivos de definição de área. + Prepares and starts a new recording session. Args: - output_folder (str): A pasta onde os arquivos serão salvos. - frame_width (int): A largura dos quadros de vídeo. - frame_height (int): A altura dos quadros de vídeo. - is_video_file (bool): Se True, pula a criação do arquivo de vídeo, - usado para análises de vídeos pré-gravados. + output_folder (str): The folder where files will be saved. + frame_width (int): The width of the video frames. + frame_height (int): The height of the video frames. + is_video_file (bool): If True, skips video file creation. Returns: - bool: True se a gravação começou com sucesso, False caso contrário. + bool: True if recording started successfully, False otherwise. """ if self.is_recording: - logging.warning("Attempted to start recording while already recording.") + log.warning("recorder.start.already_recording") return False os.makedirs(output_folder, exist_ok=True) self.base_name = os.path.basename(output_folder) + log_context = log.bind( + output_folder=output_folder, base_name=self.base_name + ) - # 1. Configura o VideoWriter, se não for uma análise de arquivo de vídeo. if not is_video_file: video_filename = os.path.join(output_folder, f"{self.base_name}.mp4") fourcc = cv2.VideoWriter_fourcc(*"mp4v") @@ -66,18 +62,11 @@ def start_recording( (frame_width, frame_height), ) if not self.video_writer.isOpened(): - logging.error( - f"Error: Could not open video writer for {video_filename}" - ) + log.error("recorder.video_writer.open_error", path=video_filename) return False else: - self.video_writer = ( - None # Garante que o writer seja None para análises de arquivos - ) + self.video_writer = None - # 2. Configura o CSVWriter para os dados de detecção. - # Este arquivo registra cada detecção com seu timestamp, número de quadro - # e coordenadas. csv_filename = os.path.join( output_folder, f"3_CoordMovimento_{self.base_name}.csv" ) @@ -89,21 +78,20 @@ def start_recording( ) self.csv_file.flush() except IOError as e: - logging.error(f"Error: Could not open CSV file {csv_filename}. {e}") + log.error("recorder.csv.open_error", path=csv_filename, exc_info=e) if self.video_writer: self.video_writer.release() return False - # 3. Salva os arquivos CSV de definição de área. self._save_area_definitions(output_folder) self.is_recording = True self.start_time = time.time() - logging.info(f"Started recording. Output folder: {output_folder}") + log_context.info("recorder.start.success") return True def stop_recording(self): - """Para a gravação e libera todos os manipuladores de arquivo.""" + """Stops the recording and releases all file handlers.""" if not self.is_recording: return @@ -117,28 +105,15 @@ def stop_recording(self): self.csv_writer = None self.is_recording = False - logging.info(f"Stopped recording for {self.base_name}.") + log.info("recorder.stop.success", base_name=self.base_name) def write_video_frame(self, frame): - """ - Escreve um único quadro no arquivo de vídeo, se o VideoWriter estiver ativo. - - Args: - frame: O quadro de vídeo (array numpy) a ser escrito. - """ + """Writes a single frame to the video file.""" if self.is_recording and self.video_writer: self.video_writer.write(frame) def write_detection_data(self, timestamp, frame_number, detections): - """ - Escreve os dados de uma ou mais detecções no arquivo CSV. - - Args: - timestamp (float): O timestamp da detecção em segundos. - frame_number (int): O número do quadro em que a detecção ocorreu. - detections (list): Uma lista de tuplas, onde cada tupla representa - uma detecção (x1, y1, x2, y2, confidence). - """ + """Writes detection data to the CSV file.""" if self.is_recording and self.csv_writer: for x1, y1, x2, y2, confidence in detections: self.csv_writer.writerow( @@ -152,24 +127,15 @@ def write_detection_data(self, timestamp, frame_number, detections): int(confidence * 100), ] ) - # Flush the buffer to ensure data is written to disk immediately, - # reducing the risk of data loss in case of a crash. self.csv_file.flush() - logging.info(f"Wrote {len(detections)} detections for frame {frame_number}") + log.debug( + "recorder.detections.wrote", + count=len(detections), + frame=frame_number, + ) def _save_area_definitions(self, folder_path): - """ - Salva as definições de área de processamento e de interesse em CSVs. - - - `1_ProcessingArea_...csv`: Salva as coordenadas do polígono que define - a área total onde a detecção de objetos é realizada. - - `2_AreasOfInterest_...csv`: Salva as coordenadas dos quadrados - (retângulos) que definem as áreas de interesse específicas. - - Args: - folder_path (str): A pasta onde os arquivos CSV serão salvos. - """ - # Salva a Área de Processamento (Polígono) + """Saves processing and interest area definitions to CSVs.""" processing_area_filename = os.path.join( folder_path, f"1_ProcessingArea_{self.base_name}.csv" ) @@ -180,7 +146,6 @@ def _save_area_definitions(self, folder_path): f.flush() os.fsync(f.fileno()) - # Salva as Áreas de Interesse (Quadrados) areas_of_interest_filename = os.path.join( folder_path, f"2_AreasOfInterest_{self.base_name}.csv" ) @@ -192,7 +157,7 @@ def _save_area_definitions(self, folder_path): f.flush() os.fsync(f.fileno()) - logging.info(f"Saved area definitions to {folder_path}") + log.info("recorder.area_definitions.saved", path=folder_path) if __name__ == "__main__": @@ -215,10 +180,9 @@ def _save_area_definitions(self, folder_path): print("\nRecording started successfully.") # Test writing data - recorder.recording_start_frame = 100 # Simulate starting mid-stream - for i in range(10): # Simulate 10 frames + recorder.recording_start_frame = 100 + for i in range(10): frame_num = 100 + i - # Add some changing element to the frame cv2.putText( dummy_frame, f"Frame {frame_num}", @@ -230,10 +194,9 @@ def _save_area_definitions(self, folder_path): ) recorder.write_video_frame(dummy_frame) - # Simulate a detection if i % 2 == 0: detections = [(100 + i, 150, 200 + i, 250, 0.95)] - recorder.write_detection_data(frame_num, detections) + recorder.write_detection_data(time.time(), frame_num, detections) time.sleep(0.1) diff --git a/src/zebtrack/io/sources.py b/src/zebtrack/io/sources.py index 3f2dec3..ff130b4 100644 --- a/src/zebtrack/io/sources.py +++ b/src/zebtrack/io/sources.py @@ -38,5 +38,6 @@ def create_source(source_type: str, **kwargs: Any) -> FrameSource: return VideoFileSource(video_path=video_path) else: raise ValueError( - f"Unsupported source type: {source_type}. Supported types are 'camera', 'file'." + f"Unsupported source type: {source_type}. " + "Supported types are 'camera', 'file'." ) diff --git a/src/zebtrack/io/video_source.py b/src/zebtrack/io/video_source.py index ea28344..6e5fc8e 100644 --- a/src/zebtrack/io/video_source.py +++ b/src/zebtrack/io/video_source.py @@ -3,35 +3,26 @@ do `cv2.VideoCapture` para lidar com arquivos de vídeo como fontes de quadros. """ -import logging import os from typing import Any, Dict, Tuple import cv2 import numpy as np +import structlog from zebtrack.io.frame_source import FrameSource +log = structlog.get_logger() + class VideoFileSource(FrameSource): """ Representa um arquivo de vídeo como uma fonte de quadros. - - Esta classe encapsula um objeto `cv2.VideoCapture` para abrir um arquivo de - vídeo, ler suas propriedades (largura, altura, FPS), e fornecer quadros - um por um. """ def __init__(self, video_path: str): """ Inicializa a fonte de vídeo a partir de um caminho de arquivo. - - Args: - video_path (str): O caminho para o arquivo de vídeo. - - Raises: - FileNotFoundError: Se o arquivo de vídeo não for encontrado. - IOError: Se o arquivo de vídeo não puder ser aberto pelo OpenCV. """ if not os.path.exists(video_path): raise FileNotFoundError(f"Video file not found at: {video_path}") @@ -42,29 +33,25 @@ def __init__(self, video_path: str): if not self.cap.isOpened(): raise IOError(f"Cannot open video file: {video_path}") - # Store video properties self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) self.fps = self.cap.get(cv2.CAP_PROP_FPS) if self.fps == 0: - print("Warning: Video FPS is 0. Defaulting to 30.") + log.warning("video.fps.zero", path=video_path) self.fps = 30 self.frame_count = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)) - logging.info(f"Video source loaded: {os.path.basename(video_path)}") - logging.info( - f"Properties: {self.width}x{self.height} @ {self.fps:.2f} FPS, " - f"{self.frame_count} frames total." + log.info( + "video.source.loaded", + path=video_path, + width=self.width, + height=self.height, + fps=self.fps, + frame_count=self.frame_count, ) def get_frame(self) -> Tuple[bool, np.ndarray | None]: - """ - Reads the next frame from the video file. - - Returns: - A tuple containing a boolean (success) and the frame (numpy array). - Returns (False, None) at the end of the video. - """ + """Reads the next frame from the video file.""" ret, frame = self.cap.read() if not ret: return False, None @@ -75,13 +62,7 @@ def get_current_frame_number(self) -> float: return self.cap.get(cv2.CAP_PROP_POS_FRAMES) def get_properties(self) -> Dict[str, Any]: - """ - Retorna um dicionário com as propriedades do vídeo. - - Returns: - dict: Um dicionário contendo largura, altura, FPS e contagem total - de quadros. - """ + """Returns a dictionary with the video properties.""" return { "width": self.width, "height": self.height, @@ -90,16 +71,15 @@ def get_properties(self) -> Dict[str, Any]: } def release(self) -> None: - """Libera o recurso do arquivo de vídeo.""" + """Releases the video file resource.""" if self.cap.isOpened(): self.cap.release() - logging.info(f"Video source released: {os.path.basename(self.video_path)}") + log.info("video.source.released", path=self.video_path) if __name__ == "__main__": print("Testing VideoFileSource...") - # Create a dummy video file for testing since we can't assume one exists. test_video_path = "test_video.mp4" frame_width, frame_height = 640, 480 fps = 30 @@ -110,10 +90,8 @@ def release(self) -> None: if not writer.isOpened(): print("Failed to create a dummy video writer.") else: - # Write 100 black frames with a frame number for i in range(100): - frame = cv2.UMat(frame_height, frame_width, cv2.CV_8UC3) - frame.setTo(0) + frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8) text = f"Frame {i + 1}" cv2.putText( frame, text, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2 @@ -122,35 +100,27 @@ def release(self) -> None: writer.release() print(f"Created a dummy video file: {test_video_path}") - # Now, test the VideoFileSource with the created file try: video_source = VideoFileSource(test_video_path) - frame_counter = 0 while True: ret, frame = video_source.get_frame() if not ret: print("\nEnd of video reached.") break - frame_counter += 1 - # To test, we can just show a few frames if frame_counter <= 5 or frame_counter >= 95: reported_frame = video_source.get_current_frame_number() print( f"Read frame number: {frame_counter} " f"(reported: {reported_frame})" ) - # cv2.imshow("Test Video", frame) # Can't show in this env - # cv2.waitKey(30) print(f"\nTotal frames read: {frame_counter}") video_source.release() - except (FileNotFoundError, IOError) as e: print(f"Error: {e}") finally: - # Clean up the dummy video file if os.path.exists(test_video_path): os.remove(test_video_path) print(f"Cleaned up dummy video file: {test_video_path}") diff --git a/src/zebtrack/plugins/__init__.py b/src/zebtrack/plugins/__init__.py new file mode 100644 index 0000000..dc81d83 --- /dev/null +++ b/src/zebtrack/plugins/__init__.py @@ -0,0 +1,9 @@ +from .openvino_detector import OpenVINOPlugin +from .yolo_detector import YOLOv8Plugin + +# A simple plugin registry. The keys are the user-facing names. +# The main application can use this to discover and instantiate plugins. +DETECTOR_PLUGINS = { + YOLOv8Plugin.get_name(): YOLOv8Plugin, + OpenVINOPlugin.get_name(): OpenVINOPlugin, +} diff --git a/src/zebtrack/plugins/base.py b/src/zebtrack/plugins/base.py new file mode 100644 index 0000000..597b28a --- /dev/null +++ b/src/zebtrack/plugins/base.py @@ -0,0 +1,54 @@ +from abc import ABC, abstractmethod +from typing import List, Tuple + +import numpy as np + + +class DetectorPlugin(ABC): + """ + Abstract Base Class for a detector plugin. + + This interface defines the contract that all detector plugins must follow, + ensuring they can be used interchangeably by the main Detector class. + """ + + @abstractmethod + def __init__(self, model_path: str): + """ + Initializes the plugin and loads the specified model. + + Args: + model_path (str): The path to the model file or directory. + """ + pass + + @abstractmethod + def detect(self, frame: np.ndarray) -> List[Tuple[int, int, int, int, float]]: + """ + Performs object detection on a single frame. + + Args: + frame (np.ndarray): The input video frame. + + Returns: + A list of detections. Each detection is a tuple containing: + (x1, y1, x2, y2, confidence). + """ + pass + + @staticmethod + @abstractmethod + def get_name() -> str: + """ + Returns the user-friendly name of the plugin. + e.g., "YOLOv8 (Ultralytics)" + """ + pass + + @property + @abstractmethod + def model_input_shape(self) -> Tuple[int, int]: + """ + Returns the expected input shape (height, width) of the model. + """ + pass diff --git a/src/zebtrack/plugins/openvino_detector.py b/src/zebtrack/plugins/openvino_detector.py new file mode 100644 index 0000000..4dc8a06 --- /dev/null +++ b/src/zebtrack/plugins/openvino_detector.py @@ -0,0 +1,133 @@ +import glob +import os +from typing import List, Tuple + +import cv2 +import numpy as np +import openvino as ov +import torch +from ultralytics.utils.ops import non_max_suppression, scale_boxes + +from zebtrack.plugins.base import DetectorPlugin +from zebtrack.settings import settings + + +class OpenVINOPlugin(DetectorPlugin): + """A detector plugin that uses an OpenVINO-optimized model.""" + + def __init__(self, model_path: str): + """ + Initializes the plugin and loads the OpenVINO model. + + Args: + model_path (str): Path to the directory containing the .xml and .bin files. + """ + self.conf_threshold = settings.yolo_model.confidence_threshold + self.nms_threshold = settings.yolo_model.nms_threshold + + xml_files = glob.glob(os.path.join(model_path, "*.xml")) + if not xml_files: + raise FileNotFoundError( + f"Could not find a .xml model file in directory: {model_path}" + ) + + model_xml_path = xml_files[0] + core = ov.Core() + model = core.read_model(model_xml_path) + self.compiled_model = core.compile_model(model=model, device_name="AUTO") + self.input_layer = self.compiled_model.input(0) + self.output_layer = self.compiled_model.output(0) + self.infer_request = self.compiled_model.create_infer_request() + + def detect(self, frame: np.ndarray) -> List[Tuple[int, int, int, int, float]]: + """Performs inference using the OpenVINO model.""" + input_tensor = self._preprocess(frame) + self.infer_request.infer({self.input_layer.any_name: input_tensor}) + results = self.infer_request.results + predictions = self._postprocess(results, frame.shape) + return predictions + + def _preprocess(self, frame: np.ndarray) -> np.ndarray: + """Prepares a frame for OpenVINO inference using letterboxing.""" + n, c, h, w = self.input_layer.shape + letterboxed_frame, _, _ = _letterbox(frame, new_shape=(w, h), auto=False) + rgb_frame = cv2.cvtColor(letterboxed_frame, cv2.COLOR_BGR2RGB) + input_tensor = rgb_frame.transpose(2, 0, 1) / 255.0 + input_tensor = np.expand_dims(input_tensor, axis=0).astype(np.float32) + return input_tensor + + def _postprocess(self, result: dict, original_frame_shape: tuple) -> list: + """Postprocesses the OpenVINO model's output.""" + output_tensor = result[self.output_layer] + preds = non_max_suppression( + prediction=torch.from_numpy(output_tensor), + conf_thres=self.conf_threshold, + iou_thres=self.nms_threshold, + agnostic=True, + ) + detections = preds[0] + if detections is None or len(detections) == 0: + return [] + + detections[:, :4] = scale_boxes( + self.model_input_shape, detections[:, :4], original_frame_shape + ).round() + + final_detections = [] + for *xyxy, conf, cls in detections: + final_detections.append( + (int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3]), float(conf)) + ) + return final_detections + + @staticmethod + def get_name() -> str: + return "OpenVINO" + + @property + def model_input_shape(self) -> Tuple[int, int]: + return self.input_layer.shape[2], self.input_layer.shape[3] # (h, w) + + +def _letterbox( + img: np.ndarray, + new_shape: tuple = (640, 640), + color: tuple = (114, 114, 114), + auto: bool = True, + scaleFill: bool = False, + scaleup: bool = True, + stride: int = 32, +): + """ + Standard letterboxing function from ultralytics. + Resizes and pads image while meeting stride-multiple constraints. + """ + shape = img.shape[:2] + if isinstance(new_shape, int): + new_shape = (new_shape, new_shape) + + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + if not scaleup: + r = min(r, 1.0) + + ratio = r, r + new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) + dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] + if auto: + dw, dh = np.mod(dw, stride), np.mod(dh, stride) + elif scaleFill: + dw, dh = 0.0, 0.0 + new_unpad = (new_shape[1], new_shape[0]) + ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] + + dw /= 2 + dh /= 2 + + if shape[::-1] != new_unpad: + img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) + top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) + left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) + img = cv2.copyMakeBorder( + img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color + ) + return img, ratio, (dw, dh) diff --git a/src/zebtrack/plugins/yolo_detector.py b/src/zebtrack/plugins/yolo_detector.py new file mode 100644 index 0000000..ca9eaf1 --- /dev/null +++ b/src/zebtrack/plugins/yolo_detector.py @@ -0,0 +1,54 @@ +from typing import List, Tuple + +import numpy as np +import torch +from ultralytics import YOLO + +from zebtrack.plugins.base import DetectorPlugin +from zebtrack.settings import settings + + +class YOLOv8Plugin(DetectorPlugin): + """A detector plugin that uses the ultralytics YOLOv8 model.""" + + def __init__(self, model_path: str): + """ + Initializes the YOLOv8 model. + + Args: + model_path (str): The path to the .pt model file. + """ + self.model = YOLO(model_path) + self.conf_threshold = settings.yolo_model.confidence_threshold + self.nms_threshold = settings.yolo_model.nms_threshold + + def detect(self, frame: np.ndarray) -> List[Tuple[int, int, int, int, float]]: + """ + Performs inference using the YOLOv8 model. + """ + use_half = torch.cuda.is_available() + results = self.model( + frame, + verbose=False, + conf=self.conf_threshold, + iou=self.nms_threshold, + half=use_half, + ) + + predictions = [] + for det in results[0].boxes.data.cpu().numpy(): + x1, y1, x2, y2, confidence, _ = det + predictions.append((int(x1), int(y1), int(x2), int(y2), float(confidence))) + + return predictions + + @staticmethod + def get_name() -> str: + return "YOLOv8 (Ultralytics)" + + @property + def model_input_shape(self) -> Tuple[int, int]: + # This is a bit of a simplification. YOLOv8 can handle various input sizes, + # but 640 is the default and what's implicitly used. + # For a more robust implementation, one might inspect the model's properties. + return (640, 640) diff --git a/src/zebtrack/settings.py b/src/zebtrack/settings.py index aec764d..44cfa00 100644 --- a/src/zebtrack/settings.py +++ b/src/zebtrack/settings.py @@ -3,13 +3,15 @@ a loader function to read and validate the configuration from a YAML file. """ -import logging from pathlib import Path from typing import List, Tuple +import structlog import yaml from pydantic import BaseModel, Field, ValidationError +log = structlog.get_logger() + # --- Pydantic Models for Configuration Structure --- @@ -30,7 +32,10 @@ class ArduinoSettings(BaseModel): port: str = Field( ..., - description="The serial port the Arduino is connected to (e.g., 'COM5' or '/dev/ttyACM0').", + description=( + "The serial port the Arduino is connected to (e.g., 'COM5' or " + "'/dev/ttyACM0')." + ), ) baud_rate: int = Field(..., description="The baud rate for serial communication.") @@ -51,20 +56,28 @@ class YOLOModelSettings(BaseModel): ..., gt=0, lt=1, - description="Non-Maximum Suppression threshold for filtering overlapping bounding boxes.", + description=( + "Non-Maximum Suppression threshold for filtering overlapping bounding " + "boxes." + ), ) class VideoProcessingSettings(BaseModel): """Settings for processing video files or live streams.""" - fps: int = Field(..., description="Frames Per Second (FPS) for saving output videos.") + fps: int = Field( + ..., description="Frames Per Second (FPS) for saving output videos." + ) processing_interval: int = Field( ..., description="Process 1 frame every N frames to optimize performance." ) processing_offset: int = Field( ..., - description="Frame offset for processing. E.g., offset=1 and interval=10 processes frames 1, 11, 21, ...", + description=( + "Frame offset for processing. E.g., offset=1 and interval=10 processes " + "frames 1, 11, 21, ..." + ), ) @@ -76,7 +89,10 @@ class DetectionZonesSettings(BaseModel): ) squares: List[Tuple[Tuple[int, int], Tuple[int, int]]] = Field( ..., - description="A list of rectangular zones, each defined by top-left and bottom-right points.", + description=( + "A list of rectangular zones, each defined by top-left and " + "bottom-right points." + ), ) colors: List[Tuple[int, int, int]] = Field( ..., description="The BGR colors for drawing each square on the overlay." @@ -96,7 +112,10 @@ class ReproducibilitySettings(BaseModel): seed: int = Field( 42, - description="Seed for random number generators (numpy, torch) to ensure consistent results.", + description=( + "Seed for random number generators (numpy, torch) to ensure consistent " + "results." + ), ) @@ -146,31 +165,31 @@ def load_settings( ValueError: If there are validation or parsing errors. """ if not default_config_path.is_file(): - logging.error(f"Default configuration file not found at: {default_config_path}") + log.error("settings.load.file_not_found", path=str(default_config_path)) raise FileNotFoundError( f"Default configuration file not found at: {default_config_path}" ) - logging.info(f"Loading base settings from {default_config_path}...") + log.info("settings.load.start", path=str(default_config_path)) try: with open(default_config_path, "r") as f: config_data = yaml.safe_load(f) if override_config_path.is_file(): - logging.info(f"Loading override settings from {override_config_path}...") + log.info("settings.load.override", path=str(override_config_path)) with open(override_config_path, "r") as f: override_data = yaml.safe_load(f) if override_data: config_data = _merge_configs(config_data, override_data) settings = Settings.model_validate(config_data) - logging.info("Settings loaded and validated successfully.") + log.info("settings.load.success") return settings except yaml.YAMLError as e: - logging.error(f"Error parsing YAML file: {e}") + log.error("settings.load.yaml_error", error=str(e)) raise ValueError(f"Error parsing YAML file: {e}") except ValidationError as e: - logging.error(f"Configuration validation error: {e}") + log.error("settings.load.validation_error", error=str(e)) raise ValueError(f"Configuration validation error: {e}") @@ -178,6 +197,6 @@ def load_settings( try: settings = load_settings() except (FileNotFoundError, ValueError) as e: - logging.critical(f"Failed to load application settings: {e}") + log.critical("settings.load.failed", error=str(e)) # In a real app, you might want to exit or use default settings settings = None diff --git a/src/zebtrack/ui/gui.py b/src/zebtrack/ui/gui.py index 742f1be..cc18f11 100644 --- a/src/zebtrack/ui/gui.py +++ b/src/zebtrack/ui/gui.py @@ -2,7 +2,6 @@ Este módulo define a interface gráfica principal (GUI) para a aplicação Zebtrack. """ -import logging import os import queue import threading @@ -23,41 +22,32 @@ ) import cv2 +import structlog # Import custom modules from zebtrack.core.detector import Detector, draw_overlay from zebtrack.io.camera import Camera from zebtrack.io.video_source import VideoFileSource +from zebtrack.plugins import DETECTOR_PLUGINS from zebtrack.settings import settings +log = structlog.get_logger() + class ApplicationGUI: """ A classe principal que gerencia a interface gráfica (a "Visão"). - - Esta classe é responsável por: - - Construir a janela principal e os widgets da interface. - - Exibir o estado da aplicação conforme ditado pelo AppController. - - Encaminhar as ações do usuário (cliques de botão, etc.) para o AppController. """ def __init__(self, root, controller): """ Inicializa a ApplicationGUI. - - Args: - root (tk.Tk): A janela raiz do Tkinter para a aplicação. - controller (AppController): A instância do controlador principal. """ self.root = root self.controller = controller self.root.title("Zebtrack Controller") self.root.protocol("WM_DELETE_WINDOW", self.controller.on_close) - # A inicialização dos módulos de backend e o gerenciamento de estado - # são de responsabilidade do AppController. - - # --- Variáveis de UI --- self.welcome_frame = None self.main_controls_frame = None self.status_var = StringVar() @@ -114,13 +104,13 @@ def _create_main_control_frame(self): self.start_rec_btn = Button( self.main_controls_frame, text="Start Recording", - command=self.controller.start_recording, # Delegated + command=self.controller.start_recording, ) self.start_rec_btn.pack(side="left", padx=5) self.stop_rec_btn = Button( self.main_controls_frame, text="Stop Recording", - command=self.controller.stop_recording, # Delegated + command=self.controller.stop_recording, state="disabled", ) self.stop_rec_btn.pack(side="left", padx=5) @@ -151,45 +141,83 @@ def _create_main_control_frame(self): ).pack(side="left", padx=10) Button( - self.main_controls_frame, text="Close Project", command=self.controller.close_project # Delegated + self.main_controls_frame, + text="Close Project", + command=self.controller.close_project, ).pack(side="left", padx=5) status_text = ( - f"Project: {self.controller.project_manager.get_project_name()} ({project_type})" + f"Project: {self.controller.project_manager.get_project_name()} " + f"({project_type})" ) self.status_var.set(status_text) Label(self.root, textvariable=self.status_var).pack(pady=5) def _load_project_view(self): """ - Transiciona da tela de boas-vindas para a visualização de controle principal. + Transitions from the welcome screen to the main control view and + initializes the detector with the appropriate plugin. """ - self.controller.detector = Detector(self.controller.project_manager) + pm = self.controller.project_manager + use_openvino = pm.project_data.get("use_openvino", False) + + try: + if use_openvino: + plugin_name = "OpenVINO" + model_path = pm.project_data.get("openvino_model_path") + if not model_path or not os.path.exists(model_path): + raise ValueError("OpenVINO model path not found or invalid.") + else: + plugin_name = "YOLOv8 (Ultralytics)" + model_path = settings.yolo_model.path + + plugin_class = DETECTOR_PLUGINS.get(plugin_name) + if not plugin_class: + raise ValueError(f"Detector plugin '{plugin_name}' not found.") + + plugin_instance = plugin_class(model_path=model_path) + self.controller.detector = Detector(plugin=plugin_instance) + + except (ValueError, FileNotFoundError) as e: + log.error("detector.init.failed", error=str(e), exc_info=True) + self.show_error( + "Detector Error", f"Failed to initialize the detector: {e}" + ) + self._create_welcome_frame() + return self._create_main_control_frame() - project_type = self.controller.project_manager.get_project_type() + project_type = pm.get_project_type() if project_type == "live": if not self.controller.arduino.connect(): - self.show_warning("Arduino Warning", "Could not connect to Arduino. Running in offline mode.") + self.show_warning( + "Arduino Warning", + "Could not connect to Arduino. Running in offline mode.", + ) try: self.controller.camera = Camera() self.controller.active_frame_source = self.controller.camera self.controller.detector.update_scaling( - self.controller.camera.actual_width, self.controller.camera.actual_height + self.controller.camera.actual_width, + self.controller.camera.actual_height, ) except IOError as e: self.show_error("Camera Error", str(e)) self._create_welcome_frame() return elif project_type == "pre-recorded": - next_video = self.controller.project_manager.get_next_video() + next_video = pm.get_next_video() if next_video is None: self.process_video_btn.config(state="disabled") - self.set_status(f"Project: {self.controller.project_manager.get_project_name()} - All videos processed.") + self.set_status( + f"Project: {pm.get_project_name()} - All videos processed." + ) else: video_name = os.path.basename(next_video) - self.set_status(f"Project: {self.controller.project_manager.get_project_name()} - Ready to process: {video_name}") + self.set_status( + f"Project: {pm.get_project_name()} - Ready to process: {video_name}" + ) if project_type == "live": self.controller.capture_thread = threading.Thread( @@ -213,7 +241,7 @@ def _live_frame_capture_loop(self): ret, frame = self.controller.active_frame_source.get_frame() if not ret: - logging.error("Capture thread: Failed to get frame from live source.") + log.error("gui.capture_thread.get_frame_failed") time.sleep(0.5) continue @@ -252,7 +280,7 @@ def _live_processing_loop(self): self.controller.on_close() break cv2.destroyAllWindows() - logging.info("Live processing loop finished and destroyed CV2 windows.") + log.info("gui.live_processing_loop.finished") def _file_processing_loop(self): """ @@ -261,7 +289,7 @@ def _file_processing_loop(self): if not self.controller.is_recording or not isinstance( self.controller.active_frame_source, VideoFileSource ): - logging.error("File processing loop started in an invalid state.") + log.error("gui.file_processing_loop.invalid_state") return show_preview = self.show_preview_var.get() @@ -297,7 +325,7 @@ def _file_processing_loop(self): break frame_number = target_frame - logging.info(f"Processing frame {frame_number}...") + log.info("gui.file_processing_loop.progress", frame=frame_number) if not show_preview and total_frames > 0: progress_percent = int((frame_number / total_frames) * 100) @@ -313,7 +341,6 @@ def _file_processing_loop(self): if show_preview: draw_overlay(frame, detections, self.controller.detector) - # ... (drawing logic remains the same) cv2.imshow("File Processing", frame) if cv2.waitKey(1) & 0xFF == ord("q"): self.controller.program_exit_event.set() @@ -368,7 +395,6 @@ def _create_project_workflow(self): use_openvino = self.use_openvino_var.get() - # Delegate the logic to the controller self.controller.create_project_workflow(project_path, project_type, use_openvino, video_files) def _open_project_workflow(self): @@ -377,22 +403,24 @@ def _open_project_workflow(self): if not project_path: return - # Delegate the logic to the controller self.controller.open_project_workflow(project_path) def _define_groups(self): """Permite que o usuário defina nomes para os grupos de tratamento.""" group_count = self.ask_string("Number of Groups", "Enter the total number of groups:") if group_count is not None: - group_names = [] - for i in range(int(group_count)): - name = self.ask_string("Group Name", f"Enter name for group {i + 1}:") - if name: - group_names.append(name) - self.controller.project_manager.project_data["groups"] = group_names - self.controller.project_manager.save_project() - self.show_info("Success", "Group names have been updated.") - + try: + num_groups = int(group_count) + group_names = [] + for i in range(num_groups): + name = self.ask_string("Group Name", f"Enter name for group {i + 1}:") + if name: + group_names.append(name) + self.controller.project_manager.project_data["groups"] = group_names + self.controller.project_manager.save_project() + self.show_info("Success", "Group names have been updated.") + except (ValueError, TypeError): + self.show_error("Invalid Input", "Please enter a valid number for the group count.") def _on_close(self): """Delegates the close action to the controller.""" @@ -402,7 +430,6 @@ def _join_threads(self): """Delegates thread joining to the controller.""" self.controller.join_threads() - # --- View Helper Methods --- def set_status(self, text): """Updates the UI status bar.""" self.status_var.set(text) @@ -437,11 +464,11 @@ def ask_open_filenames(self, title, filetypes): def update_button_state(self, button_name, state): """Updates the state of a button ('normal' or 'disabled').""" - if button_name == "start_rec": + if button_name == "start_rec" and hasattr(self, 'start_rec_btn'): self.start_rec_btn.config(state=state) - elif button_name == "stop_rec": + elif button_name == "stop_rec" and hasattr(self, 'stop_rec_btn'): self.stop_rec_btn.config(state=state) - elif button_name == "process_video": + elif button_name == "process_video" and hasattr(self, 'process_video_btn'): self.process_video_btn.config(state=state) def ask_recording_details(self, group_names): @@ -473,5 +500,6 @@ def on_confirm(): if __name__ == "__main__": + # Using print is fine here as it's for direct execution feedback print("This file is intended to be imported, not run directly.") - print("Run main.py to start the application.") + print("Run the main application script to start Zebtrack.") diff --git a/src/zebtrack/utils.py b/src/zebtrack/utils.py index c51641c..5c596df 100644 --- a/src/zebtrack/utils.py +++ b/src/zebtrack/utils.py @@ -1,9 +1,11 @@ -import logging import random import numpy as np +import structlog import torch +log = structlog.get_logger() + def set_seed(seed: int): """ @@ -24,4 +26,4 @@ def set_seed(seed: int): # consistency is more important than performance. torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False - logging.info(f"Set random seed to {seed} for reproducibility.") + log.info("reproducibility.seed.set", seed=seed) diff --git a/tests/__pycache__/__init__.cpython-312.pyc b/tests/__pycache__/__init__.cpython-312.pyc index 3c932b32138017ddac9fdb8a4863a6384df1c77f..a7a8b3d9eb25902fe43ed976e0043c3b128e61ac 100644 GIT binary patch delta 21 bcmXRf=Q_>H%f$c$S6)u!@@Bj>F-#EvL16}N delta 21 bcmXRf=Q_>H%f$c$Y1bxlc{Bc+7^VmSJn9Aw diff --git a/tests/__pycache__/test_project_manager.cpython-312-pytest-8.4.1.pyc b/tests/__pycache__/test_project_manager.cpython-312-pytest-8.4.1.pyc index dac70fa4fe9a8825c6b5499eccd64b11004c4d3a..4b8bdbeb61d81b6985c0168e44b03f1feeefe5fd 100644 GIT binary patch delta 115 zcmdmBGuMXeG%qg~0}xz!xsi*LneoNy$sRIc+X SU*uKZJWoKI7ijV+*)#xLNFkU2 delta 128 zcmbPhv%!YzG%qg~0}!NL+sMVq%=l}v2(uL%6HuU$VRIPsHa4cG43i7_WhZNMrE*;6 zb@{-+!f7~pGM7Iukp2$hd;}4j-*7dvFse;1;Ge*GnOFG(11qO1<78%mdI1pQJDB?s QOm6NL(B=hNvt2d~08ULIZU6uP diff --git a/tests/__pycache__/test_recorder.cpython-312-pytest-8.4.1.pyc b/tests/__pycache__/test_recorder.cpython-312-pytest-8.4.1.pyc index b03a38bb17a97cc307c58b61cb9d58065833c1b4..f9a3939c5916e2a424e9855cd746e02279085cdc 100644 GIT binary patch delta 183 zcmdn%I@y)$G%qg~0}xz!xshvyFcUk|XOFU(B5eBarDnm+I`aPs!^ zcJf~5P`SvVa+O2vyERbwBCqzZ?_xmh7kG8S44n(S+M92RPGw=_o18DT3v8Z|v==8( W?R8#lpz_T<(nr`BZ%qzV$^if>$u}+TC$^V5-I4|>Re_&wc3}T!dFRC~>T*R6M z%$X`;!hM-n7b3MmaB`xk{N%SH7HlBNF9MSpMT|J0I{APa^_W>X`6f>g6;T4oe*~NN z>$@0;12F;2;RCVxHt!Lg%EBbg2y}t|A N`Uo52ugU64IRKV=HU$6x diff --git a/tests/test_controller.py b/tests/test_controller.py index dbf92af..420c3c4 100644 --- a/tests/test_controller.py +++ b/tests/test_controller.py @@ -3,6 +3,7 @@ from src.zebtrack.core.controller import AppController + class TestAppController(unittest.TestCase): @patch('src.zebtrack.core.controller.Arduino') diff --git a/tests/test_detector.py b/tests/test_detector.py index cd1681c..016e832 100644 --- a/tests/test_detector.py +++ b/tests/test_detector.py @@ -4,168 +4,119 @@ import numpy as np from zebtrack.core.detector import Detector +from zebtrack.plugins.base import DetectorPlugin from zebtrack.settings import settings +class MockDetectorPlugin(DetectorPlugin): + """A mock plugin for testing the main Detector class.""" + + def __init__(self, model_path: str = "mock_model"): + self.model_path = model_path + self._detect_return_value = [] + + def detect(self, frame: np.ndarray): + # Allow configuring the return value for different test cases + return self._detect_return_value + + @staticmethod + def get_name() -> str: + return "Mock Plugin" + + @property + def model_input_shape(self): + return (640, 480) + + # Test helper to configure the mock's output + def set_detect_return_value(self, value): + self._detect_return_value = value + + class TestDetector(unittest.TestCase): def setUp(self): - """Set up a detector instance for tests.""" - self.yolo_patcher = patch("zebtrack.core.detector.YOLO") - self.openvino_patcher = patch("zebtrack.core.detector.ov") - self.glob_patcher = patch("zebtrack.core.detector.glob") - - self.mock_yolo = self.yolo_patcher.start() - self.mock_openvino = self.openvino_patcher.start() - self.mock_glob = self.glob_patcher.start() - - # Configure a default mock project manager to prevent OpenVINO loading - self.mock_project_manager = MagicMock() - self.mock_project_manager.project_data = {"use_openvino": False, "openvino_model_path": ""} - - self.detector = Detector(project_manager=self.mock_project_manager) - - def tearDown(self): - """Stop the patchers.""" - self.yolo_patcher.stop() - self.openvino_patcher.stop() - self.glob_patcher.stop() - - def test_initialization_default_yolo(self): - """Test that the detector initializes with YOLO by default.""" - # The setUp method already configures this scenario - self.assertIsNotNone(self.detector.model) - self.assertFalse(self.detector.is_openvino) - self.mock_yolo.assert_called_with(settings.yolo_model.path) - - def test_initialization_openvino(self): - """Test that the detector initializes with OpenVINO when configured.""" - self.mock_project_manager.project_data = { - "use_openvino": True, - "openvino_model_path": "/fake/path/model.xml", - } - # Mock glob to find a fake model file - self.mock_glob.glob.return_value = ["/fake/path/model.xml"] - - detector = Detector(project_manager=self.mock_project_manager) - - self.assertTrue(detector.is_openvino) - self.assertIsNotNone(detector.compiled_model) - self.mock_openvino.Core.assert_called() + """Set up a detector instance with a mock plugin for tests.""" + self.mock_plugin = MockDetectorPlugin() + self.detector = Detector(plugin=self.mock_plugin) + + def test_initialization(self): + """Test that the detector initializes correctly with a plugin.""" + self.assertIsNotNone(self.detector) + self.assertEqual(self.detector.plugin, self.mock_plugin) + self.assertEqual(self.detector.flag, 0) + + def test_initialization_fails_without_plugin(self): + """Test that Detector raises an error if no plugin is provided.""" + with self.assertRaises(ValueError): + Detector(plugin=None) def test_update_scaling(self): """Test the logic for scaling detection zones.""" base_width = settings.camera.desired_width base_height = settings.camera.desired_height - - test_width = 640 - test_height = 360 - + test_width, test_height = 640, 360 self.detector.update_scaling(test_width, test_height) - scale_x = test_width / base_width scale_y = test_height / base_height - original_point = self.detector.base_polygon[0] scaled_point = self.detector.scaled_polygon[0] - expected_x = int(original_point[0] * scale_x) expected_y = int(original_point[1] * scale_y) - self.assertEqual(scaled_point[0], expected_x) self.assertEqual(scaled_point[1], expected_y) - def test_process_frame_yolo_path(self): - """Test the frame processing logic using the YOLO path.""" - mock_results = MagicMock() - fake_detection = np.array([[10, 10, 50, 50, 0.9, 0]]) - mock_results[0].boxes.data.cpu.return_value.numpy.return_value = fake_detection - self.detector.model.return_value = mock_results - - dummy_frame = np.zeros((settings.camera.desired_height, settings.camera.desired_width, 3), dtype=np.uint8) - - with patch.object(self.detector, '_is_inside_polygon', return_value=True): - detections, command = self.detector.process_frame(dummy_frame, "live") - - self.assertEqual(len(detections), 1) - self.assertEqual(detections[0], (10, 10, 50, 50, 0.9)) - self.assertIsNone(command) - - def test_initialization_with_model_path_override(self): - """Test that model_path argument overrides the settings.""" - override_path = "some/other/model.pt" - detector = Detector(project_manager=self.mock_project_manager, model_path=override_path) - self.mock_yolo.assert_called_with(override_path) - - def test_initialization_yolo_load_fails(self): - """Test graceful failure when YOLO model loading raises an exception.""" - self.mock_yolo.side_effect = Exception("Model file is corrupted") - # The __init__ should catch the exception and not crash - detector = Detector(project_manager=self.mock_project_manager) - self.assertIsNone(detector.model) + def test_process_frame_delegates_to_plugin(self): + """Test that process_frame calls the plugin's detect method.""" + dummy_frame = np.zeros((480, 640, 3), dtype=np.uint8) + self.mock_plugin.detect = MagicMock(return_value=[]) + self.detector.process_frame(dummy_frame, "live") + self.mock_plugin.detect.assert_called_once_with(dummy_frame) def test_is_inside_square(self): """Test the _is_inside_square helper method.""" square = ((100, 100), (200, 200)) - # Bbox completely inside self.assertTrue(self.detector._is_inside_square(110, 110, 190, 190, square)) - # Bbox overlapping self.assertTrue(self.detector._is_inside_square(150, 150, 250, 250, square)) - # Bbox outside self.assertFalse(self.detector._is_inside_square(300, 300, 400, 400, square)) - # Bbox touching edge - self.assertTrue(self.detector._is_inside_square(90, 90, 100, 100, square)) def test_is_inside_polygon(self): """Test the _is_inside_polygon helper method.""" - # A simple square polygon for testing - polygon = np.array([[100, 100], [200, 100], [200, 200], [100, 200]], dtype=np.int32) - # Point inside + polygon = np.array([[100, 100], [200, 100], [200, 200], [100, 200]]) self.assertTrue(self.detector._is_inside_polygon(150, 150, 160, 160, polygon)) - # Point outside self.assertFalse(self.detector._is_inside_polygon(300, 300, 310, 310, polygon)) - # Point on edge - self.assertTrue(self.detector._is_inside_polygon(100, 100, 110, 110, polygon)) def test_state_machine_logic(self): """Test the command generation logic based on state.""" # Setup: A detection inside the first configured square - square = settings.detection_zones.squares[0] # e.g., ((150, 490), (360, 660)) + square = settings.detection_zones.squares[0] x_c = (square[0][0] + square[1][0]) // 2 y_c = (square[0][1] + square[1][1]) // 2 - - mock_results = MagicMock() - # A detection right in the middle of the first square - fake_detection = np.array([[x_c - 5, y_c - 5, x_c + 5, y_c + 5, 0.9, 0]]) - mock_results[0].boxes.data.cpu.return_value.numpy.return_value = fake_detection - self.detector.model.return_value = mock_results - dummy_frame = np.zeros((settings.camera.desired_height, settings.camera.desired_width, 3), dtype=np.uint8) # --- Step 1: Object enters a square, should generate ENTER command --- - # Ensure polygon check passes for this test + fake_detection = [(x_c - 5, y_c - 5, x_c + 5, y_c + 5, 0.9)] + self.mock_plugin.set_detect_return_value(fake_detection) + with patch.object(self.detector, '_is_inside_polygon', return_value=True): detections, command = self.detector.process_frame(dummy_frame, "live") self.assertEqual(self.detector.flag, 1, "Flag should be 1 (waiting for exit)") - self.assertEqual(self.detector.current_square, 1, "Should register entering square 1") + self.assertEqual(self.detector.current_square, 1) self.assertEqual(command, settings.detection_zones.enter_commands[0]) - # --- Step 2: Object is still inside a square, should generate NO command --- + # --- Step 2: Object is still inside, should generate NO command --- with patch.object(self.detector, '_is_inside_polygon', return_value=True): detections, command = self.detector.process_frame(dummy_frame, "live") - self.assertIsNone(command, "No command should be sent if object is still inside") # --- Step 3: Object moves outside all squares, should generate EXIT command --- - # New detection is outside all squares - mock_results[0].boxes.data.cpu.return_value.numpy.return_value = np.array([[10, 10, 20, 20, 0.9, 0]]) - + self.mock_plugin.set_detect_return_value([(10, 10, 20, 20, 0.9)]) with patch.object(self.detector, '_is_inside_polygon', return_value=True): detections, command = self.detector.process_frame(dummy_frame, "live") - self.assertEqual(self.detector.flag, 0, "Flag should reset to 0 (waiting for entry)") - self.assertEqual(self.detector.current_square, 0, "Current square should be reset") + self.assertEqual(self.detector.flag, 0, "Flag should reset to 0") + self.assertEqual(self.detector.current_square, 0) self.assertEqual(command, settings.detection_zones.exit_commands[0]) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_recorder.py b/tests/test_recorder.py index 2fb626d..ff8a64a 100644 --- a/tests/test_recorder.py +++ b/tests/test_recorder.py @@ -1,7 +1,6 @@ import csv import os import shutil -import sys import unittest import numpy as np diff --git a/tests/test_settings.py b/tests/test_settings.py index d974f36..de57d00 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -1,8 +1,6 @@ import unittest from unittest.mock import mock_open, patch -import yaml - from zebtrack.settings import Settings, load_settings diff --git a/tests/test_sources.py b/tests/test_sources.py index b471313..ea89edd 100644 --- a/tests/test_sources.py +++ b/tests/test_sources.py @@ -115,7 +115,7 @@ def test_create_source_file_wrong_kwarg_type(): # --- Tests for Camera class --- import time -from unittest.mock import MagicMock, PropertyMock +from unittest.mock import MagicMock from zebtrack.io import Camera From 32e18ad0b780609c1343e78ec3053f64ec7f775b Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sat, 16 Aug 2025 03:01:46 +0000 Subject: [PATCH 2/2] fix: Resolve CI failures and finalize linting This commit addresses the CI failures from the previous submission by fixing the workflow and completing all linting corrections. The `ModuleNotFoundError` in the test job was resolved by removing the `--no-root` flag from the `poetry install` step in the `ci.yml` workflow. This ensures the `zebtrack` package is properly installed in the CI environment. Additionally, all remaining linting errors, primarily line-length issues (E501), have been fixed across the codebase. The project now passes all `ruff` checks. --- .github/workflows/ci.yml | 4 +- src/zebtrack/core/project_manager.py | 10 +-- src/zebtrack/settings.py | 8 ++- src/zebtrack/ui/gui.py | 64 ++++++++++++++---- ...test_recorder.cpython-312-pytest-8.4.1.pyc | Bin 8851 -> 8835 bytes tests/test_controller.py | 6 +- tests/test_detector.py | 11 ++- tests/test_sources.py | 9 +-- 8 files changed, 78 insertions(+), 34 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bb462d1..36c32e8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -20,7 +20,7 @@ jobs: pip install poetry - name: Install dependencies run: | - poetry install --no-root + poetry install - name: Run ruff run: | poetry run ruff check . @@ -38,7 +38,7 @@ jobs: pip install poetry - name: Install dependencies run: | - poetry install --no-root + poetry install - name: Run tests run: | poetry run pytest diff --git a/src/zebtrack/core/project_manager.py b/src/zebtrack/core/project_manager.py index 800f567..0cb820d 100644 --- a/src/zebtrack/core/project_manager.py +++ b/src/zebtrack/core/project_manager.py @@ -122,7 +122,8 @@ def create_new_project( "OpenVINO Export Error", ( "Failed to move the exported OpenVINO model to the " - f"cache directory.\nPlease check permissions.\n\nError: {move_exc}" + "cache directory.\nPlease check permissions.\n\n" + f"Error: {move_exc}" ), ) return False @@ -131,9 +132,10 @@ def create_new_project( messagebox.showerror( "OpenVINO Export Error", ( - "An unexpected error occurred during OpenVINO model export.\n" - "Ensure all dependencies are installed correctly and the model " - f"path is valid.\n\nError: {e}" + "An unexpected error occurred during OpenVINO model " + "export.\nEnsure all dependencies are installed " + "correctly and the model path is valid.\n\n" + f"Error: {e}" ), ) return False diff --git a/src/zebtrack/settings.py b/src/zebtrack/settings.py index 44cfa00..cdad8ac 100644 --- a/src/zebtrack/settings.py +++ b/src/zebtrack/settings.py @@ -99,11 +99,15 @@ class DetectionZonesSettings(BaseModel): ) enter_commands: List[int] = Field( ..., - description="List of commands to send to Arduino when an object enters a square.", + description=( + "List of commands to send to Arduino when an object enters a square." + ), ) exit_commands: List[int] = Field( ..., - description="List of commands to send to Arduino when an object exits a square.", + description=( + "List of commands to send to Arduino when an object exits a square." + ), ) diff --git a/src/zebtrack/ui/gui.py b/src/zebtrack/ui/gui.py index cc18f11..46e2ffa 100644 --- a/src/zebtrack/ui/gui.py +++ b/src/zebtrack/ui/gui.py @@ -249,7 +249,10 @@ def _live_frame_capture_loop(self): if not self.controller.frame_queue.full(): self.controller.frame_queue.put((live_frame_count, frame.copy())) - if self.controller.is_capturing_for_video and not self.controller.video_queue.full(): + if ( + self.controller.is_capturing_for_video + and not self.controller.video_queue.full() + ): self.controller.video_queue.put(frame.copy()) time.sleep(1 / (settings.video_processing.fps * 1.5)) @@ -265,7 +268,9 @@ def _live_processing_loop(self): continue if self.controller.is_processing: - detections, command = self.controller.detector.process_frame(frame, "live") + detections, command = self.controller.detector.process_frame( + frame, "live" + ) if command is not None: self.controller.arduino.send_command(command) if self.controller.is_recording and detections: @@ -305,7 +310,10 @@ def _file_processing_loop(self): total_frames = video_source.get_properties()["frame_count"] frame_number = -1 - while not self.controller.program_exit_event.is_set() and frame_number < total_frames: + while ( + not self.controller.program_exit_event.is_set() + and frame_number < total_frames + ): target_frame = ( ( settings.video_processing.processing_offset @@ -329,15 +337,21 @@ def _file_processing_loop(self): if not show_preview and total_frames > 0: progress_percent = int((frame_number / total_frames) * 100) - video_name = os.path.basename(self.controller.currently_processing_video) + video_name = os.path.basename( + self.controller.currently_processing_video + ) status_msg = f"Processing: {video_name} ({progress_percent}%)" self.root.after(0, self.status_var.set, status_msg) - detections, _ = self.controller.detector.process_frame(frame, "pre-recorded") + detections, _ = self.controller.detector.process_frame( + frame, "pre-recorded" + ) if detections: props = video_source.get_properties() timestamp = frame_number / props["fps"] if props["fps"] > 0 else 0 - self.controller.recorder.write_detection_data(timestamp, frame_number, detections) + self.controller.recorder.write_detection_data( + timestamp, frame_number, detections + ) if show_preview: draw_overlay(frame, detections, self.controller.detector) @@ -357,13 +371,18 @@ def _create_project_workflow(self): if not base_path: return - project_name = self.ask_string("Project Name", "Enter a name for the new project:") + project_name = self.ask_string( + "Project Name", "Enter a name for the new project:" + ) if not project_name: return project_path = os.path.join(base_path, project_name) if os.path.exists(project_path) and os.listdir(project_path): - self.show_error("Error", "A project folder with this name already exists and is not empty.") + self.show_error( + "Error", + "A project folder with this name already exists and is not empty.", + ) return type_window = Toplevel(self.root) @@ -376,10 +395,14 @@ def _create_project_workflow(self): variable=self.use_openvino_var, ).pack(padx=20, pady=5) Button( - type_window, text="Live Analysis", command=lambda: [type_var.set("live"), type_window.destroy()] + type_window, + text="Live Analysis", + command=lambda: [type_var.set("live"), type_window.destroy()], ).pack(fill="x", padx=20, pady=5) Button( - type_window, text="Pre-recorded Analysis", command=lambda: [type_var.set("pre-recorded"), type_window.destroy()] + type_window, + text="Pre-recorded Analysis", + command=lambda: [type_var.set("pre-recorded"), type_window.destroy()], ).pack(fill="x", padx=20, pady=5) self.root.wait_window(type_window) project_type = type_var.get() @@ -389,13 +412,18 @@ def _create_project_workflow(self): video_files = [] if project_type == "pre-recorded": - video_files = self.ask_open_filenames(title="Select Video Files", filetypes=[("Video files", "*.mp4 *.avi")]) + video_files = self.ask_open_filenames( + title="Select Video Files", + filetypes=[("Video files", "*.mp4 *.avi")], + ) if not video_files: return use_openvino = self.use_openvino_var.get() - self.controller.create_project_workflow(project_path, project_type, use_openvino, video_files) + self.controller.create_project_workflow( + project_path, project_type, use_openvino, video_files + ) def _open_project_workflow(self): """Handles the UI part of opening a project, then calls the controller.""" @@ -407,20 +435,26 @@ def _open_project_workflow(self): def _define_groups(self): """Permite que o usuário defina nomes para os grupos de tratamento.""" - group_count = self.ask_string("Number of Groups", "Enter the total number of groups:") + group_count = self.ask_string( + "Number of Groups", "Enter the total number of groups:" + ) if group_count is not None: try: num_groups = int(group_count) group_names = [] for i in range(num_groups): - name = self.ask_string("Group Name", f"Enter name for group {i + 1}:") + name = self.ask_string( + "Group Name", f"Enter name for group {i + 1}:" + ) if name: group_names.append(name) self.controller.project_manager.project_data["groups"] = group_names self.controller.project_manager.save_project() self.show_info("Success", "Group names have been updated.") except (ValueError, TypeError): - self.show_error("Invalid Input", "Please enter a valid number for the group count.") + self.show_error( + "Invalid Input", "Please enter a valid number for the group count." + ) def _on_close(self): """Delegates the close action to the controller.""" diff --git a/tests/__pycache__/test_recorder.cpython-312-pytest-8.4.1.pyc b/tests/__pycache__/test_recorder.cpython-312-pytest-8.4.1.pyc index f9a3939c5916e2a424e9855cd746e02279085cdc..3b3cd11324007201ee56125370f0940c5de5f921 100644 GIT binary patch delta 191 zcmbR2+U&}EnwOW00SFwv%+Ju9$ScXXZKArHI(G^~3R4bGE^ibsBLk4loWqyPAH~ne zz{KFrkiyc!kiweEnWZo}PE1NDN-$L@g)LiyfuTr79CmoxV; zGKy_xVHILyl-?}PJ&l9WeDgZtd2Eb{n|-AsSs80KFPBecWaOF5t7y+CGTB!#dGc1p m8u==qlNf=x*cV8AU}j`wywAYU!qCihn}PE-1MlQuB?kcWB`t^m delta 205 zcmZp6o$ShcnwOW00SK{6viCxC>|iql*60L7sbcO0Aw@g z@aGCd2{1A+F}O3Nu(U9wu%>cNUdSghc@HaxP?S)ra0*+t1_MKp5EDZxSC+Zs$<9g+00yNo AUH||9 diff --git a/tests/test_controller.py b/tests/test_controller.py index 420c3c4..c7e7381 100644 --- a/tests/test_controller.py +++ b/tests/test_controller.py @@ -73,8 +73,10 @@ def test_create_project_workflow_failure(self): ) # --- Assert --- - self.mock_view.show_error.assert_called_once_with("Error", "Failed to create the new project.") + self.mock_view.show_error.assert_called_once_with( + "Error", "Failed to create the new project." + ) self.mock_view._load_project_view.assert_not_called() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/test_detector.py b/tests/test_detector.py index 016e832..fbf5cc4 100644 --- a/tests/test_detector.py +++ b/tests/test_detector.py @@ -90,7 +90,10 @@ def test_state_machine_logic(self): square = settings.detection_zones.squares[0] x_c = (square[0][0] + square[1][0]) // 2 y_c = (square[0][1] + square[1][1]) // 2 - dummy_frame = np.zeros((settings.camera.desired_height, settings.camera.desired_width, 3), dtype=np.uint8) + dummy_frame = np.zeros( + (settings.camera.desired_height, settings.camera.desired_width, 3), + dtype=np.uint8, + ) # --- Step 1: Object enters a square, should generate ENTER command --- fake_detection = [(x_c - 5, y_c - 5, x_c + 5, y_c + 5, 0.9)] @@ -104,9 +107,11 @@ def test_state_machine_logic(self): self.assertEqual(command, settings.detection_zones.enter_commands[0]) # --- Step 2: Object is still inside, should generate NO command --- - with patch.object(self.detector, '_is_inside_polygon', return_value=True): + with patch.object(self.detector, "_is_inside_polygon", return_value=True): detections, command = self.detector.process_frame(dummy_frame, "live") - self.assertIsNone(command, "No command should be sent if object is still inside") + self.assertIsNone( + command, "No command should be sent if object is still inside" + ) # --- Step 3: Object moves outside all squares, should generate EXIT command --- self.mock_plugin.set_detect_return_value([(10, 10, 20, 20, 0.9)]) diff --git a/tests/test_sources.py b/tests/test_sources.py index ea89edd..dbb5d04 100644 --- a/tests/test_sources.py +++ b/tests/test_sources.py @@ -1,10 +1,12 @@ import os +import time +from unittest.mock import MagicMock import cv2 import numpy as np import pytest -from zebtrack.io import FrameSource, VideoFileSource, create_source +from zebtrack.io import Camera, FrameSource, VideoFileSource, create_source @pytest.fixture @@ -114,11 +116,6 @@ def test_create_source_file_wrong_kwarg_type(): # --- Tests for Camera class --- -import time -from unittest.mock import MagicMock - -from zebtrack.io import Camera - @pytest.fixture def mock_video_capture(monkeypatch):