Crab detection tab

.gitignore

@@ -40,4 +40,7 @@ eeprom.bin
 *.DS_Store
 
 # Mavproxy Files
-*.parm
+*.parm
+
+# Training Dataset
+src/training/dataset/*

src/surface/gui/gui/operator_app.py

@@ -7,12 +7,14 @@
 from gui.widgets.heartbeat import HeartbeatWidget
 from gui.widgets.ip_widget import IPWidget
 from gui.widgets.logger import Logger
+from gui.widgets.tabs.crab_detector_tab import CrabDetectorTab
 from gui.widgets.tabs.general_debug_tab import GeneralDebugTab
 from gui.widgets.tabs.shipwreck import ShipwreckTab
 from gui.widgets.temperature import TemperatureSensor
 from gui.widgets.timer import InteractiveTimer
 
 SHIPWRECK_TEXT = 'Shipwreck'
+CRAB_DETECTOR_TEXT = 'Crab Detector'
 
 
 class OperatorApp(App):
@@ -55,6 +57,8 @@ def __init__(self) -> None:
         self.tabs.addTab(GeneralDebugTab(), 'General Debug')
         self.shipwreck_tab = ShipwreckTab()
         self.tabs.addTab(self.shipwreck_tab, SHIPWRECK_TEXT)
+        self.crab_detector_tab = CrabDetectorTab()
+        self.tabs.addTab(self.crab_detector_tab, CRAB_DETECTOR_TEXT)
         self.tabs.currentChanged.connect(self.changed_tabs)
         root_layout.addWidget(self.tabs)
 
@@ -65,6 +69,9 @@ def tab_change_slot(self, index: int) -> None:
         if self.tabs.tabText(index) == SHIPWRECK_TEXT:
             # Allow keyboard events
             self.shipwreck_tab.setFocus(Qt.FocusReason.TabFocusReason)
+        elif self.tabs.tabText(index) == CRAB_DETECTOR_TEXT:
+            # Allow keyboard events
+            self.crab_detector_tab.setFocus(Qt.FocusReason.TabFocusReason)
 
 
 def run_gui_operator() -> None:

src/surface/gui/gui/widgets/tabs/crab_detector_tab.py

@@ -0,0 +1,194 @@
+from dataclasses import dataclass
+from enum import Enum, IntEnum
+from typing import Generic, TypeGuard, TypeVar, override
+
+from PyQt6.QtCore import Qt, pyqtSlot
+from PyQt6.QtGui import QColor
+from PyQt6.QtWidgets import (
+    QHBoxLayout,
+    QTabWidget,
+    QVBoxLayout,
+    QWidget,
+)
+
+from gui.widgets.video_widget import (
+    CameraDescription,
+    CameraType,
+    PauseableVideoWidget,
+)
+from rov_msgs.msg import Intrinsics
+
+# CAM0_TOPIC = 'cam0/image_raw'
+CAM1_TOPIC = 'cam1/image_raw'
+
+FRAME_WIDTH = 816
+FRAME_HEIGHT = 510
+# FRAME_WIDTH = 1280
+# FRAME_HEIGHT = 800
+
+ZOOMED_WIDGET_SIZE = 405
+ZOOM_DEFAULT_IDX = 2
+ZOOMED_VIEWPORT_SIZES = (27, 45, 81, 135, 405)  # Odd factors of 405
+
+LENGTH_SCALE_FACTOR = 1.34
+
+PADDING = 200
+
+POINTS_PER_EYE = 2
+
+BASELINE_MM = 60.6
+TUBE_RADIUS_MM = 40
+
+DIVISION_SAFETY = 0.0001
+
+BLACK = QColor(Qt.GlobalColor.black)
+
+SHIPWRECK_BOW_LENGTH_CM = 30 + 16.6
+
+
+class Eye(IntEnum):
+    LEFT = 0
+    RIGHT = 1
+
+
+class Crosshair(Enum):
+    Empty = 0
+    Dot = 1
+
+
+KEYS_TO_POINT_IDX = {
+    Qt.Key.Key_1: (Eye.LEFT, 0),
+    Qt.Key.Key_2: (Eye.LEFT, 1),
+    Qt.Key.Key_3: (Eye.RIGHT, 0),
+    Qt.Key.Key_4: (Eye.RIGHT, 1),
+}
+
+
+T = TypeVar('T', int, float)
+
+
+@dataclass
+class Point2D(Generic[T]):
+    x: T
+    y: T
+
+    @override
+    def __str__(self) -> str:
+        return f'({round(self.x, 3)}, {round(self.y, 3)})'
+
+
+@dataclass
+class Point3D:
+    x: float
+    y: float
+    z: float
+
+    @override
+    def __str__(self) -> str:
+        return f'({round(self.x, 3)}, {round(self.y, 3)}, {round(self.z, 3)})'
+
+
+def has_all_points(
+    key_points: dict[Eye, list[Point2D[int] | None]],
+) -> TypeGuard['dict[Eye, list[Point2D[int]]]']:
+    return all(
+        len(key_points[eye]) == POINTS_PER_EYE
+        and all(point is not None for point in key_points[eye])
+        for eye in Eye
+    )
+
+
+def format_length(length: float) -> str:
+    return f'{length}, {length + SHIPWRECK_BOW_LENGTH_CM}'
+
+
+class CrabDetectorTab(QWidget):
+    # click_left_signal = pyqtSignal(QMouseEvent)
+    # click_right_signal = pyqtSignal(QMouseEvent)
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.img_points: dict[Eye, list[Point2D[int] | None]] = {
+            Eye.LEFT: [None, None],
+            Eye.RIGHT: [None, None],
+        }
+
+        self.keys: dict[int, bool] = {
+            Qt.Key.Key_1.value: False,
+            Qt.Key.Key_2.value: False,
+            Qt.Key.Key_3.value: False,
+            Qt.Key.Key_4.value: False,
+        }
+
+        self.crosshair = Crosshair.Dot
+        self.viewport_zoom_level = ZOOM_DEFAULT_IDX
+
+        tabs = QTabWidget()
+        tabs.addTab(self.make_coarse_tab(), 'Coarse')
+        # tabs.addTab(self.make_fine_tab(), 'Fine')
+
+        root_layout = QVBoxLayout()
+        root_layout.addWidget(tabs)
+        self.setLayout(root_layout)
+
+        # Make sure we can get keyboard
+        self.setFocusPolicy(Qt.FocusPolicy.StrongFocus)
+
+    def make_coarse_tab(self) -> QWidget:
+        # self.click_left_signal.connect(self.click_left_slot)
+        # self.click_right_signal.connect(self.click_right_slot)
+
+        cam_layout = QHBoxLayout()
+
+        # TODO: RESET THIS TO ACTUAL LUXONIS CAM ONCE IT'S WORKING AGAIN
+        self.eye_widgets = {
+            Eye.LEFT: PauseableVideoWidget(
+                CameraDescription(
+                    CameraType.ETHERNET,
+                    CAM1_TOPIC,
+                    'Down Camera',
+                    FRAME_WIDTH,
+                    FRAME_HEIGHT,
+                ),
+                #'switch_rect_stream',
+                # make_label=lambda: ClickableLabel(self.click_left_signal),
+            ),
+        }
+
+        for eye_widget in self.eye_widgets.values():
+            cam_layout.addWidget(eye_widget)
+
+        coarse_tab = QWidget()
+        coarse_tab.setLayout(cam_layout)
+
+        return coarse_tab
+
+    @staticmethod
+    def px_to_mm(px: float) -> float:
+        # 3 um/px (https://docs.luxonis.com/hardware/sensors/OV9782)
+        # / 1000 to get mm
+        return px * 3 / 1000
+
+    @pyqtSlot(Intrinsics)
+    def intrinsics_left_slot(self, intrinsics: Intrinsics) -> None:
+        self.intrinsics_left = intrinsics
+        self.show_intrinsics()
+
+    @pyqtSlot(Intrinsics)
+    def intrinsics_right_slot(self, intrinsics: Intrinsics) -> None:
+        self.intrinsics_right = intrinsics
+        self.show_intrinsics()
+
+    def show_intrinsics(self) -> None:
+        if self.intrinsics_left is None or self.intrinsics_right is None:
+            return
+
+        focal_left_mm = Point2D(
+            CrabDetectorTab.px_to_mm(self.intrinsics_left.fx),
+            CrabDetectorTab.px_to_mm(self.intrinsics_left.fy),
+        )
+        focal_right_mm = Point2D(
+            CrabDetectorTab.px_to_mm(self.intrinsics_right.fx),
+            CrabDetectorTab.px_to_mm(self.intrinsics_right.fy),
+        )

src/surface/gui/gui/widgets/video_widget.py

@@ -1,6 +1,7 @@
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from enum import IntEnum
+from pathlib import Path
 from typing import NamedTuple
 
 import cv2
@@ -12,6 +13,7 @@
 from PyQt6.QtWidgets import QLabel, QPushButton, QVBoxLayout, QWidget
 from rclpy.qos import qos_profile_default
 from sensor_msgs.msg import Image
+from ultralytics import YOLO
 
 from gui.gui_node import GUINode
 from rov_msgs.msg import VideoWidgetSwitch
@@ -309,13 +311,29 @@ def __init__(self, camera_description: CameraDescription) -> None:
             GUINode().get_logger().error('Missing Layout')
 
         self.is_paused = False
+        self.ran_model = False
 
     @pyqtSlot(Image)
     def handle_frame(self, frame: Image) -> None:
         if not self.is_paused:
             super().handle_frame(frame)
+        elif self.is_paused and not self.ran_model:
+            # cv_image = super().get_cv_image(frame)
+            cv_image = self.cv_bridge.imgmsg_to_cv2(frame, desired_encoding='passthrough')
+
+            if self.camera_description.type == CameraType.ETHERNET:
+                # Switches ethernet's color profile from BayerBGR to BGR
+                cv_image = cv2.cvtColor(cv_image, cv2.COLOR_BAYER_BGGR2BGR)
+            # Run model on cv_image
+            model = YOLO(str(Path('/home/rov/rov-26/src/surface/gui/gui/best.pt')))
+
+            results = model(cv_image)
+            results[0].show()
+            self.ran_model = True
 
     def toggle(self) -> None:
         """Toggle whether this widget is paused or playing."""
+        if self.is_paused:
+            self.ran_model = False
         self.is_paused = not self.is_paused
         self.button.setText(self.PAUSED_TEXT if self.is_paused else self.PLAYING_TEXT)

src/surface/gui/package.xml

@@ -17,6 +17,7 @@
   <exec_depend>python3-pyqtgraph</exec_depend>
   <exec_depend>pyqt6-dev-tools</exec_depend>
   <exec_depend>python3-pyqt6.qtmultimedia</exec_depend>
+  <exec_depend>python3-ultralytics-pip</exec_depend>
 
   <exec_depend>rov_msgs</exec_depend>
 

src/training/YOLO.py

@@ -0,0 +1,22 @@
+from ultralytics import YOLO
+
+# Load a pretrained YOLO26n model
+model = YOLO('yolo26n.pt')
+
+# Train the model on the COCO8 dataset for 100 epochs
+train_results = model.train(
+    data='config.yaml',  # Path to dataset configuration file
+    epochs=50,  # Number of training epochs
+    imgsz=640,  # Image size for training
+    device='mps',  # Device to run on (e.g., 'cpu', 0, [0,1,2,3])
+)
+
+# Evaluate the model's performance on the validation set
+metrics = model.val()
+
+# Perform object detection on an image
+results = model('Test Image.png')  # Predict on an image
+results[0].show()  # Display results
+
+# Export the model to ONNX format for deployment
+path = model.export(format='onnx')  # Returns the path to the exported model

src/training/config.yaml

@@ -0,0 +1,5 @@
+path: dataset
+train: images/train
+val: images/val
+nc: 1
+names: ['European Green Crab']

src/training/predict.py

@@ -0,0 +1,6 @@
+from ultralytics import YOLO
+
+model = YOLO('runs/detect/train4/weights/best.pt')
+
+results = model('dataset/images/test/frame_470.jpg')
+results[0].show()