def __init__(self, camera_arm_control_service):
self._gripper_detector = GripperDetector()
self._playfield_detector = PlayfieldDetector()
self._treasure_pickup_detector = TreasurePickupDetector()
self._treasure_pickup_success_detector = TreasurePickupSuccessDetector()
self._end_zone_treasure_detector = EndZoneTreasureDetector()
self._camera_arm_control_service = camera_arm_control_service
class LocalWorldDetector(object):
MAX_IMAGE_SIZE = Vector2(600, 800)
def __init__(self, camera_arm_control_service):
self._gripper_detector = GripperDetector()
self._playfield_detector = PlayfieldDetector()
self._treasure_pickup_detector = TreasurePickupDetector()
self._treasure_pickup_success_detector = TreasurePickupSuccessDetector()
self._end_zone_treasure_detector = EndZoneTreasureDetector()
self._camera_arm_control_service = camera_arm_control_service
def detect_objects(self, image, detection_mode, local_world_map):
self._resize_image(image)
local_world_descriptor = LocalWorldDescriptor()
local_world_descriptor.reference_image = image
if detection_mode != LocalWorldMapDetectionMode.NONE:
self._execute_detection(image, detection_mode, local_world_descriptor, local_world_map)
return local_world_descriptor
def _execute_detection(self, image, detection_mode, local_world_descriptor, local_world_map):
try:
self._execute_detection_mode(image, detection_mode, local_world_descriptor, local_world_map)
except DetectionError as e:
Logger.get_instance().log(self, "Local World Detection Failed", LogEntryType.WARNING, error = e, blob_data = image)
def _resize_image(self, image):
if image.size > self.MAX_IMAGE_SIZE:
image.resize(image.size.shrink_to_fit(self.MAX_IMAGE_SIZE))
def _execute_detection_mode(self, image, detection_mode, local_world_descriptor, local_world_map):
if detection_mode == LocalWorldMapDetectionMode.TREASURE_PICKUP_DETECTION:
self._detect_pickup_treasures(image, local_world_descriptor)
elif detection_mode == LocalWorldMapDetectionMode.TREASURE_PICKUP_SUCCESS_DETECTION:
self._detect_treasure_pickup_success(image, local_world_descriptor, local_world_map)
elif detection_mode == LocalWorldMapDetectionMode.END_ZONE_TREASURE_DETECTION:
self._detect_end_zone_treasures(image, local_world_descriptor)
def _detect_pickup_treasures(self, image, local_world_descriptor):
local_world_descriptor.local_world_object_types = [LocalWorldObjectType.GRIPPER, LocalWorldObjectType.TREASURE]
self._camera_arm_control_service.move_camera(TreasurePickupDetector.DETECTION_CAMERA_ANGLE)
if self._camera_arm_control_service.movement_in_progress:
return
local_world_descriptor.gripper = self._gripper_detector.detect(image)
if local_world_descriptor.gripper is not None:
local_world_descriptor.treasures = self._treasure_pickup_detector.detect(image, local_world_descriptor.gripper)
if len(local_world_descriptor.treasures) == 0:
raise DetectionError("No treasure detected.")
def _detect_treasure_pickup_success(self, image, local_world_descriptor, local_world_map):
local_world_descriptor.local_world_object_types = [LocalWorldObjectType.TREASURE]
self._camera_arm_control_service.move_camera(TreasurePickupSuccessDetector.DETECTION_CAMERA_ANGLE)
if self._camera_arm_control_service.movement_in_progress:
return
if local_world_map.contains(LocalWorldObjectType.GRIPPER):
local_world_descriptor.treasures = self._treasure_pickup_success_detector.detect(image, local_world_map.gripper)
else:
raise DetectionError("Gripper must be detected prior performing this operation.")
def _detect_end_zone_treasures(self, image, local_world_descriptor):
local_world_descriptor.local_world_object_types = [LocalWorldObjectType.PLAYFIELD, LocalWorldObjectType.TREASURE]
self._camera_arm_control_service.move_camera_sequenced(EndZoneTreasureDetector.DETECTION_CAMERA_MOVEMENT_SEQUENCE)
if self._camera_arm_control_service.movement_in_progress:
return
try:
local_world_descriptor.playfield = self._playfield_detector.detect(image)
if local_world_descriptor.playfield is not None:
local_world_descriptor.treasures = self._end_zone_treasure_detector.detect(image, local_world_descriptor.playfield)
except DetectionError as e:
self._camera_arm_control_service.move_next()
raise e