diff --git a/autonomy/lib/constants.dart b/autonomy/lib/constants.dart index b6e121fc..a9a60470 100644 --- a/autonomy/lib/constants.dart +++ b/autonomy/lib/constants.dart @@ -48,6 +48,28 @@ class Constants { /// The maximum time to spend searching for an aruco tag static const Duration arucoSearchTimeout = Duration(seconds: 20); + /// The maximum time to spend searching for a vision object + static const Duration visionSearchTimeout = Duration(seconds: 20); + /// The camera that should be used to detect Aruco tags static const CameraName arucoDetectionCamera = CameraName.ROVER_FRONT; + + /// The size of the search area in meters for the lawnmower pattern + static const double searchAreaMeters = 20; + + /// The width of each strip in the lawnmower pattern in meters + /// TODO: calibrate based on camera detection range testing in lab + static const double searchStripWidthMeters = 3; + + /// How close (in degrees) the detected object yaw must be to center + static const double visionCenterYawEpsilon = 3; + + /// How close (in pixels) the detected object center must be to frame center + static const int visionCenterPixelEpsilon = 20; + + /// How close (normalized -1..1) the detected object x position must be + static const double visionCenterXPositionEpsilon = 0.1; + + /// How many consecutive frames the target can be missing before we stop + static const int visionMissingFramesToArrive = 3; } diff --git a/autonomy/lib/src/drive/drive_interface.dart b/autonomy/lib/src/drive/drive_interface.dart index c87ca3ba..dce47df3 100644 --- a/autonomy/lib/src/drive/drive_interface.dart +++ b/autonomy/lib/src/drive/drive_interface.dart @@ -139,4 +139,10 @@ abstract class DriveInterface extends Service { /// Drive forward to approach an Aruco tag Future approachAruco() async {} + + /// Spin to face a detected object, returns whether or not it was found + Future spinForObject( + List types, { + CameraName? desiredCamera, + }) async => false; } diff --git a/autonomy/lib/src/drive/rover_drive.dart b/autonomy/lib/src/drive/rover_drive.dart index 91064713..fa5538e5 100644 --- a/autonomy/lib/src/drive/rover_drive.dart +++ b/autonomy/lib/src/drive/rover_drive.dart @@ -83,6 +83,12 @@ class RoverDrive extends DriveInterface { @override Future approachAruco() => sensorDrive.approachAruco(); + @override + Future spinForObject( + List types, { + CameraName? desiredCamera, + }) => sensorDrive.spinForObject(types, desiredCamera: desiredCamera); + @override Future faceOrientation(Orientation orientation) async { if (useImu) { diff --git a/autonomy/lib/src/drive/sensor_drive.dart b/autonomy/lib/src/drive/sensor_drive.dart index 3b5302ff..66dda33e 100644 --- a/autonomy/lib/src/drive/sensor_drive.dart +++ b/autonomy/lib/src/drive/sensor_drive.dart @@ -175,6 +175,36 @@ class SensorDrive extends DriveInterface with RoverDriveCommands { return foundAruco; } + @override + Future spinForObject( + List types, { + CameraName? desiredCamera, + }) async { + setThrottle(config.turnThrottle); + var foundObject = true; + foundObject = + await runFeedback(() { + if (!foundObject) return true; + spinLeft(); + return types.any( + (type) => + collection.video.getDetection( + type, + desiredCamera: desiredCamera, + ) != + null, + ); + }).timeout( + Constants.visionSearchTimeout, + onTimeout: () { + foundObject = false; + return false; + }, + ); + await stop(); + return foundObject; + } + @override Future approachAruco() async { // const sizeThreshold = 0.2; diff --git a/autonomy/lib/src/orchestrator/orchestrator_interface.dart b/autonomy/lib/src/orchestrator/orchestrator_interface.dart index 1c2a2918..3780a70e 100644 --- a/autonomy/lib/src/orchestrator/orchestrator_interface.dart +++ b/autonomy/lib/src/orchestrator/orchestrator_interface.dart @@ -44,6 +44,12 @@ abstract class OrchestratorInterface extends Service { handleGpsTask(command); case AutonomyTask.VISUAL_MARKER: handleArucoTask(command); + case AutonomyTask.HAMMER_TARGET: + handleHammerTask(command); + case AutonomyTask.ROCK_HAMMER_TARGET: + handleHammerTask(command); + case AutonomyTask.BOTTLE_TARGET: + handleBottleTask(command); } } diff --git a/autonomy/lib/src/orchestrator/rover_orchestrator.dart b/autonomy/lib/src/orchestrator/rover_orchestrator.dart index 8e05835d..7c4e466a 100644 --- a/autonomy/lib/src/orchestrator/rover_orchestrator.dart +++ b/autonomy/lib/src/orchestrator/rover_orchestrator.dart @@ -1,8 +1,9 @@ +import "dart:async"; import "dart:math"; import "package:autonomy/constants.dart"; import "package:autonomy/interfaces.dart"; -import "dart:async"; +import "package:autonomy/src/state_machine/rover_states/approach_visible_object.dart"; import "package:coordinate_converter/coordinate_converter.dart"; @@ -198,6 +199,32 @@ class RoverOrchestrator extends OrchestratorInterface with ValueReporter { return true; } + /// Generates GPS waypoints for a lawnmower search pattern centered on [center]. + /// + /// Covers a [Constants.searchAreaMeters] x [Constants.searchAreaMeters] area, + /// sweeping back and forth in strips of [Constants.searchStripWidthMeters] width. + List generateLawnmowerWaypoints(GpsCoordinates center) { + final waypoints = []; + final centerUtm = center.toUTM(); + const halfSize = Constants.searchAreaMeters / 2; + const stripWidth = Constants.searchStripWidthMeters; + + var row = 0; + for (var y = -halfSize; y <= halfSize; y += stripWidth) { + final xStart = row.isEven ? -halfSize : halfSize; + final xEnd = row.isEven ? halfSize : -halfSize; + waypoints.add( + (centerUtm + UTMCoordinates(x: xStart, y: y, zoneNumber: 1)).toGps(), + ); + waypoints.add( + (centerUtm + UTMCoordinates(x: xEnd, y: y, zoneNumber: 1)).toGps(), + ); + row++; + } + + return waypoints; + } + @override void handleGpsTask(AutonomyCommand command) { final destination = command.destination; @@ -291,7 +318,7 @@ class RoverOrchestrator extends OrchestratorInterface with ValueReporter { currentState = AutonomyState.SEARCHING; collection.logger.info("Searching for ArUco tag"); - final didSeeAruco = await collection.drive.spinForAruco( + await collection.drive.spinForAruco( command.arucoId, desiredCamera: Constants.arucoDetectionCamera, ); @@ -300,12 +327,37 @@ class RoverOrchestrator extends OrchestratorInterface with ValueReporter { desiredCamera: Constants.arucoDetectionCamera, ); - if (!didSeeAruco || detectedAruco == null) { + if (detectedAruco == null) { + collection.logger.info( + "ArUco not found after spin, starting lawnmower search", + ); + final waypoints = generateLawnmowerWaypoints(collection.gps.coordinates); + for (final waypoint in waypoints) { + if (currentCommand == null) return; + await calculateAndFollowPath( + waypoint, + abortOnError: false, + alternateEndCondition: () => + collection.video.getArucoDetection( + command.arucoId, + desiredCamera: Constants.arucoDetectionCamera, + ) != + null, + ); + detectedAruco = collection.video.getArucoDetection( + command.arucoId, + desiredCamera: Constants.arucoDetectionCamera, + ); + if (detectedAruco != null) break; + } + } + + if (detectedAruco == null) { collection.logger.error("Could not find desired Aruco tag"); currentState = AutonomyState.NO_SOLUTION; currentCommand = null; return; - } + } collection.logger.info("Found aruco"); currentState = AutonomyState.APPROACHING; @@ -425,8 +477,133 @@ class RoverOrchestrator extends OrchestratorInterface with ValueReporter { } @override - Future handleHammerTask(AutonomyCommand command) async {} + Future handleHammerTask(AutonomyCommand command) async { + await _handleVisualObjectTask( + command: command, + targetTypes: [ + DetectedObjectType.HAMMER, + DetectedObjectType.ROCK_HAMMER, + ], + label: "hammer", + ); + } @override - Future handleBottleTask(AutonomyCommand command) async {} + Future handleBottleTask(AutonomyCommand command) async { + await _handleVisualObjectTask( + command: command, + targetTypes: [DetectedObjectType.BOTTLE], + label: "bottle", + ); + } + + Future _handleVisualObjectTask({ + required AutonomyCommand command, + required List targetTypes, + required String label, + }) async { + collection.drive.setLedStrip(ProtoColor.RED); + + if (command.destination != GpsCoordinates(latitude: 0, longitude: 0)) { + if (!await calculateAndFollowPath( + command.destination, + abortOnError: false, + )) { + collection.logger.error( + "Failed to follow path towards initial destination", + ); + currentState = AutonomyState.NO_SOLUTION; + currentCommand = null; + return; + } + } + + currentState = AutonomyState.SEARCHING; + collection.logger.info("Searching for $label"); + await collection.drive.spinForObject( + targetTypes, + desiredCamera: Constants.arucoDetectionCamera, + ); + var detection = _firstTargetDetection(targetTypes); + + if (detection == null) { + collection.logger.info( + "$label not found after spin, starting lawnmower search", + ); + final waypoints = generateLawnmowerWaypoints(collection.gps.coordinates); + for (final waypoint in waypoints) { + if (currentCommand == null) return; + final reached = await calculateAndFollowPath( + waypoint, + abortOnError: false, + alternateEndCondition: () => _firstTargetDetection(targetTypes) != null, + ); + detection = _firstTargetDetection(targetTypes); + if (detection != null) break; + if (!reached) { + await collection.drive.spinForObject( + targetTypes, + desiredCamera: Constants.arucoDetectionCamera, + ); + detection = _firstTargetDetection(targetTypes); + if (detection != null) break; + } + } + } + + if (detection == null) { + collection.logger.error("Could not find $label"); + currentState = AutonomyState.NO_SOLUTION; + currentCommand = null; + return; + } + + collection.logger.info("Found $label, approaching"); + currentState = AutonomyState.APPROACHING; + + controller.pushState( + ApproachVisibleObject( + controller, + collection: collection, + targetTypes: targetTypes, + desiredCamera: Constants.arucoDetectionCamera, + ), + ); + + executionTimer = PeriodicTimer(const Duration(milliseconds: 10), ( + timer, + ) async { + if (currentCommand == null) { + onCommandEnd(); + timer.cancel(); + return; + } + + if (!controller.hasState()) { + collection.logger.info("Reached target $label"); + collection.drive.setLedStrip(ProtoColor.GREEN, blink: true); + currentState = AutonomyState.AT_DESTINATION; + onCommandEnd(); + timer.cancel(); + return; + } + + controller.update(); + }); + } + + DetectedObjectSnapshot? _firstTargetDetection( + List targetTypes, + ) { + for (final type in targetTypes) { + final detection = collection.video.getDetection( + type, + desiredCamera: Constants.arucoDetectionCamera, + ); + if (detection != null) { + return detection; + } + } + return null; + } } diff --git a/autonomy/lib/src/state_machine/rover_states/approach_visible_object.dart b/autonomy/lib/src/state_machine/rover_states/approach_visible_object.dart new file mode 100644 index 00000000..38df312c --- /dev/null +++ b/autonomy/lib/src/state_machine/rover_states/approach_visible_object.dart @@ -0,0 +1,110 @@ +import "package:autonomy/constants.dart"; +import "package:autonomy/interfaces.dart"; +import "package:autonomy/rover.dart"; +import "package:autonomy/src/drive/drive_config.dart"; + +class ApproachVisibleObject extends RoverState { + final AutonomyInterface collection; + final List targetTypes; + final CameraName? desiredCamera; + final int? arucoId; + + DriveConfig get config => collection.drive.config; + + int _missingFrames = 0; + + ApproachVisibleObject( + super.controller, { + required this.collection, + required List targetTypes, + this.desiredCamera, + this.arucoId, + }) : targetTypes = targetTypes.isEmpty + ? [DetectedObjectType.DETECTION_TYPE_UNDEFINED] + : targetTypes; + + @override + void update() { + if (collection.drive is! RoverDrive) { + controller.popState(); + return; + } + + final drive = (collection.drive as RoverDrive).sensorDrive; + + DetectedObjectSnapshot? detection; + for (final type in targetTypes) { + detection = collection.video.getDetection( + type, + desiredCamera: desiredCamera, + arucoId: arucoId, + ); + if (detection != null) { + break; + } + } + + if (detection == null) { + _missingFrames += 1; + drive.stopMotors(); + if (_missingFrames >= Constants.visionMissingFramesToArrive) { + controller.popState(); + } + return; + } + + _missingFrames = 0; + + final object = detection.object; + final xPositionError = object.hasXPosition() ? object.xPosition : null; + final yawError = object.hasYaw() ? object.yaw : null; + + double? pixelError; + if (yawError == null && object.hasCenterX()) { + final width = detection.details.hasStreamWidth() + ? detection.details.streamWidth + : detection.details.resolutionWidth; + if (width > 0) { + pixelError = object.centerX - (width / 2.0); + } + } + + if (xPositionError == null && yawError == null && pixelError == null) { + collection.logger.warning( + "Detected object lacks centering data", + body: "Cannot determine yaw or center for ${object.objectType}", + ); + drive.stopMotors(); + return; + } + + final isCentered = xPositionError != null + ? xPositionError.abs() <= Constants.visionCenterXPositionEpsilon + : (yawError != null + ? yawError.abs() <= Constants.visionCenterYawEpsilon + : pixelError!.abs() <= Constants.visionCenterPixelEpsilon); + + if (!isCentered) { + drive.setThrottle(config.turnThrottle); + final error = xPositionError ?? yawError ?? pixelError!; + if (error > 0) { + drive.spinRight(); + } else { + drive.spinLeft(); + } + return; + } + + drive.setThrottle(config.forwardThrottle); + drive.moveForward(); + } + + @override + void exit() { + if (collection.drive is! RoverDrive) { + return; + } + final drive = (collection.drive as RoverDrive).sensorDrive; + drive.stopMotors(); + } +} diff --git a/autonomy/lib/src/video/rover_video.dart b/autonomy/lib/src/video/rover_video.dart index e3363e44..5006ede8 100644 --- a/autonomy/lib/src/video/rover_video.dart +++ b/autonomy/lib/src/video/rover_video.dart @@ -21,20 +21,38 @@ class RoverVideo extends VideoInterface { } @override - DetectedObject? getArucoDetection(int id, {CameraName? desiredCamera}) { + DetectedObjectSnapshot? getDetection( + DetectedObjectType type, { + CameraName? desiredCamera, + int? arucoId, + }) { for (final result in _cachedResults.where( (e) => e.details.name == (desiredCamera ?? e.details.name), )) { for (final object in result.detectedObjects) { - if (object.objectType == DetectedObjectType.ARUCO && - object.arucoTagId == id) { - return object; + if (object.objectType != type) { + continue; + } + if (type == DetectedObjectType.ARUCO && + arucoId != null && + object.arucoTagId != arucoId) { + continue; } + return DetectedObjectSnapshot(object: object, details: result.details); } } return null; } + @override + DetectedObject? getArucoDetection(int id, {CameraName? desiredCamera}) => + getDetection( + DetectedObjectType.ARUCO, + desiredCamera: desiredCamera, + arucoId: id, + ) + ?.object; + @override Future waitForAruco( int id, { diff --git a/autonomy/lib/src/video/video_interface.dart b/autonomy/lib/src/video/video_interface.dart index 4344054d..46d0a853 100644 --- a/autonomy/lib/src/video/video_interface.dart +++ b/autonomy/lib/src/video/video_interface.dart @@ -3,6 +3,13 @@ import "dart:async"; import "package:autonomy/constants.dart"; import "package:autonomy/interfaces.dart"; +class DetectedObjectSnapshot { + final DetectedObject object; + final CameraDetails details; + + const DetectedObjectSnapshot({required this.object, required this.details}); +} + /// Handles obstacle detection data and ArUco data from video abstract class VideoInterface extends Service with Receiver { final AutonomyInterface collection; @@ -27,6 +34,13 @@ abstract class VideoInterface extends Service with Receiver { void updateFrame(VideoData result); + DetectedObjectSnapshot? getDetection( + DetectedObjectType type, { + CameraName? desiredCamera, + int? arucoId, + }) => + null; + DetectedObject? getArucoDetection(int id, {CameraName? desiredCamera}) => null; diff --git a/burt_network/lib/src/generated/autonomy.pbenum.dart b/burt_network/lib/src/generated/autonomy.pbenum.dart index 37517594..e43940f5 100644 --- a/burt_network/lib/src/generated/autonomy.pbenum.dart +++ b/burt_network/lib/src/generated/autonomy.pbenum.dart @@ -60,16 +60,25 @@ class AutonomyTask extends $pb.ProtobufEnum { AutonomyTask._(2, _omitEnumNames ? '' : 'VISUAL_MARKER'); static const AutonomyTask BETWEEN_GATES = AutonomyTask._(3, _omitEnumNames ? '' : 'BETWEEN_GATES'); + static const AutonomyTask HAMMER_TARGET = + AutonomyTask._(4, _omitEnumNames ? '' : 'HAMMER_TARGET'); + static const AutonomyTask ROCK_HAMMER_TARGET = + AutonomyTask._(5, _omitEnumNames ? '' : 'ROCK_HAMMER_TARGET'); + static const AutonomyTask BOTTLE_TARGET = + AutonomyTask._(6, _omitEnumNames ? '' : 'BOTTLE_TARGET'); static const $core.List values = [ AUTONOMY_TASK_UNDEFINED, GPS_ONLY, VISUAL_MARKER, BETWEEN_GATES, + HAMMER_TARGET, + ROCK_HAMMER_TARGET, + BOTTLE_TARGET, ]; static final $core.List _byValue = - $pb.ProtobufEnum.$_initByValueList(values, 3); + $pb.ProtobufEnum.$_initByValueList(values, 6); static AutonomyTask? valueOf($core.int value) => value < 0 || value >= _byValue.length ? null : _byValue[value]; diff --git a/burt_network/lib/src/generated/autonomy.pbjson.dart b/burt_network/lib/src/generated/autonomy.pbjson.dart index 641c8cb6..d9cbd8ed 100644 --- a/burt_network/lib/src/generated/autonomy.pbjson.dart +++ b/burt_network/lib/src/generated/autonomy.pbjson.dart @@ -44,13 +44,17 @@ const AutonomyTask$json = { {'1': 'GPS_ONLY', '2': 1}, {'1': 'VISUAL_MARKER', '2': 2}, {'1': 'BETWEEN_GATES', '2': 3}, + {'1': 'HAMMER_TARGET', '2': 4}, + {'1': 'ROCK_HAMMER_TARGET', '2': 5}, + {'1': 'BOTTLE_TARGET', '2': 6}, ], }; /// Descriptor for `AutonomyTask`. Decode as a `google.protobuf.EnumDescriptorProto`. final $typed_data.Uint8List autonomyTaskDescriptor = $convert.base64Decode( 'CgxBdXRvbm9teVRhc2sSGwoXQVVUT05PTVlfVEFTS19VTkRFRklORUQQABIMCghHUFNfT05MWR' - 'ABEhEKDVZJU1VBTF9NQVJLRVIQAhIRCg1CRVRXRUVOX0dBVEVTEAM='); + 'ABEhEKDVZJU1VBTF9NQVJLRVIQAhIRCg1CRVRXRUVOX0dBVEVTEAMSEQoNSEFNTUVSX1RBUkdF' + 'VBAEEhYKElJPQ0tfSEFNTUVSX1RBUkdFVBAFEhEKDUJPVFRMRV9UQVJHRVQQBg=='); @$core.Deprecated('Use autonomyDataDescriptor instead') const AutonomyData$json = { diff --git a/burt_network/lib/src/serial/firmware.dart b/burt_network/lib/src/serial/firmware.dart index 2bd7d078..1f44e3c8 100644 --- a/burt_network/lib/src/serial/firmware.dart +++ b/burt_network/lib/src/serial/firmware.dart @@ -13,7 +13,8 @@ class BurtFirmwareSerial extends Service { /// The interval to read serial data at. static const readInterval = Duration(milliseconds: 10); /// How long it should take for a firmware device to respond to a handshake. - static const handshakeDelay = Duration(milliseconds: 200); + static const handshakeDelay = Duration(milliseconds: 1200); + static const handshakeTimeout = Duration(milliseconds: 1500); /// The reset code to send to a firmware device. static final resetCode = Uint8List.fromList([0, 0, 0, 0]); @@ -67,7 +68,18 @@ class BurtFirmwareSerial extends Service { // Execute the handshake if (!(await _reset() || await _reset())) logger.warning("The Teensy on port $port failed to reset"); - if (!await _sendHandshake()) { + + // Retry handshake up to 3 times with brief backoff + var handshakeSuccess = false; + for (var attempt = 0; attempt < 3; attempt++) { + if (await _sendHandshake()) { + handshakeSuccess = true; + break; + } + await Future.delayed(const Duration(milliseconds: 200)); + } + + if (!handshakeSuccess) { logger.warning("Could not connect to Teensy", body: "Device on port $port failed the handshake"); return false; } @@ -82,12 +94,32 @@ class BurtFirmwareSerial extends Service { logger.debug("Sending handshake to port $port..."); final handshake = Connect(sender: Device.SUBSYSTEMS, receiver: Device.FIRMWARE); _serial.write(handshake.writeToBuffer()); + + // Wait for MCU boot, then poll for response await Future.delayed(handshakeDelay); - final response = _serial.readBytes(4); + + final deadline = DateTime.now().add(handshakeTimeout); + final response = []; + + while (DateTime.now().isBefore(deadline) && response.length < 4) { + final chunk = _serial.readBytes(4 - response.length); + if (chunk.isNotEmpty) { + response.addAll(chunk); + } else { + await Future.delayed(const Duration(milliseconds: 20)); + } + } + if (response.isEmpty) { logger.trace("Device did not respond"); return false; } + + if (response.length < 4) { + logger.trace("Device responded with incomplete handshake: $response"); + return false; + } + try { final message = Connect.fromBuffer(response); logger.trace("Device responded with: ${message.toProto3Json()}"); diff --git a/vision/lib/analyzer.py b/vision/lib/analyzer.py index 50d02bdb..16ba784b 100644 --- a/vision/lib/analyzer.py +++ b/vision/lib/analyzer.py @@ -2,6 +2,8 @@ from ultralytics import YOLO import torch +from lib.network.src.generated import vision_pb2 + class ImageAnalyzer: def __init__(self, load_from="trained-model.pt"): @@ -15,32 +17,71 @@ def __init__(self, load_from="trained-model.pt"): self.model = YOLO(load_from, task="detect") def has_mallet(self, frame: cv2.Mat, confidence=0.80) -> tuple[bool, float]: - results = self.model.predict(frame, stream=True, verbose=False) - for result in results: - json = result.summary() - if not json: - return False, 0.0 - object = json[0] - obj_confidence = object["confidence"] - is_mallet = ("mallet" in object["name"].lower()) and obj_confidence > confidence - if is_mallet: - return True, obj_confidence - else: - return False, obj_confidence - else: - return False, 0.0 + _, detections = self.analyze_frame(frame, confidence=confidence) + for detection in detections: + if detection.object_type == vision_pb2.HAMMER: + return True, confidence + return False, 0.0 def annotateImage(self, frame: cv2.Mat, confidence=0.80) -> cv2.Mat: - results = self.model.predict(frame, stream=True, verbose=False) - for result in results: - json = result.summary() - if not json: - return frame - object = json[0] - is_mallet = object["name"] == "mallet" and object["confidence"] > confidence - if is_mallet: - return result.plot() - return frame + annotated, _ = self.analyze_frame(frame, confidence=confidence) + return annotated + + def analyze_frame(self, frame: cv2.Mat, confidence=0.80): + results = self.model.predict(frame, verbose=False) + if not results: + return frame, [] + + result = results[0] + detections = [] + + boxes = getattr(result, "boxes", None) + if boxes is None or boxes.xyxy is None: + return frame, [] + + names = self.model.names + height, width = frame.shape[:2] + + for idx in range(len(boxes)): + conf = float(boxes.conf[idx]) + if conf < confidence: + continue + + cls_id = int(boxes.cls[idx]) + label = str(names.get(cls_id, "")).lower() + object_type = self._map_label_to_object_type(label) + if object_type is None: + continue + + x1, y1, x2, y2 = boxes.xyxy[idx].tolist() + center_x = int((x1 + x2) / 2.0) + center_y = int((y1 + y2) / 2.0) + area = max(0.0, (x2 - x1) * (y2 - y1)) + + x_position = 0.0 + if width > 0: + x_position = (center_x / width) * 2.0 - 1.0 + + detection = vision_pb2.DetectedObject( + object_type=object_type, + x_position=x_position, + relative_size=(area / (width * height)) if width > 0 and height > 0 else 0.0, + center_x=center_x, + center_y=center_y, + ) + detections.append(detection) + + annotated = result.plot() if hasattr(result, "plot") else frame + return annotated, detections + + def _map_label_to_object_type(self, label: str): + if "hammer" in label and "rock" in label: + return vision_pb2.ROCK_HAMMER + if "hammer" in label or "mallet" in label: + return vision_pb2.HAMMER + if "bottle" in label: + return vision_pb2.BOTTLE + return None def getFrames(self, cam, show=False): ret, frame = cam.read() diff --git a/vision/lib/analyzer_server.py b/vision/lib/analyzer_server.py index 0a2ac1ca..c0db2d7c 100644 --- a/vision/lib/analyzer_server.py +++ b/vision/lib/analyzer_server.py @@ -35,7 +35,7 @@ def handle_frame(self, data): if self.show == "before": cv2.imshow("Rover Vision Server (before annotation)", matrix) - frame = self.analyzer.annotateImage(matrix) + frame, detections = self.analyzer.analyze_frame(matrix) if self.show == "after": cv2.imshow("Rover Vision Server (after annotation)", frame) @@ -45,7 +45,11 @@ def handle_frame(self, data): ) if not success: return - data = VideoData(frame=bytes(jpg), details=data.details) + data = VideoData( + frame=bytes(jpg), + details=data.details, + detected_objects=detections, + ) jpg_size = len(jpg) if jpg_size > self.max_packet_size: