webots 0.8.0

Webots bindings for Rust
Documentation
use std::{
    ffi::{CStr, CString},
    slice::from_raw_parts,
};

use thiserror::Error;
use webots_bindings::{
    wb_camera_recognition_disable, wb_camera_recognition_disable_segmentation,
    wb_camera_recognition_enable, wb_camera_recognition_enable_segmentation,
    wb_camera_recognition_get_number_of_objects, wb_camera_recognition_get_objects,
    wb_camera_recognition_get_sampling_period, wb_camera_recognition_get_segmentation_image,
    wb_camera_recognition_has_segmentation, wb_camera_recognition_is_segmentation_enabled,
    wb_camera_recognition_save_segmentation_image, wb_device_get_node_type, WbDeviceTag,
    WbNodeType_WB_NODE_CAMERA,
};

use crate::Camera;

#[derive(Debug, Error)]
pub enum RecognitionError {
    #[error("failed to get objects: objects data is NULL")]
    ObjectsDataIsNull,
}

pub struct RecognitionObject<'a> {
    pub id: i32,
    pub position: [f64; 3],
    pub orientation: [f64; 4],
    pub size: [f64; 2],
    pub position_on_image: [i32; 2],
    pub size_on_image: [i32; 2],
    pub number_of_colors: i32,
    pub colors: &'a [f64],
    pub model: &'a str,
}

pub struct Recognition(WbDeviceTag);

impl Recognition {
    pub(crate) fn new(camera_device: WbDeviceTag) -> Self {
        assert_eq!(WbNodeType_WB_NODE_CAMERA, unsafe {
            wb_device_get_node_type(camera_device)
        });
        Self(camera_device)
    }

    pub fn enable(&self, sampling_period: i32) {
        unsafe { wb_camera_recognition_enable(self.0, sampling_period) }
    }

    pub fn disable(&self) {
        unsafe { wb_camera_recognition_disable(self.0) }
    }

    pub fn get_sampling_period(&self) -> i32 {
        unsafe { wb_camera_recognition_get_sampling_period(self.0) }
    }

    pub fn get_number_of_objects(&self) -> i32 {
        unsafe { wb_camera_recognition_get_number_of_objects(self.0) }
    }

    pub fn get_objects<'a>(&self) -> Result<Vec<RecognitionObject<'a>>, RecognitionError> {
        let number_of_objects = self.get_number_of_objects();
        let objects = unsafe {
            let objects = wb_camera_recognition_get_objects(self.0);
            if objects.is_null() {
                return Err(RecognitionError::ObjectsDataIsNull);
            }
            from_raw_parts(objects, number_of_objects as usize)
        };
        Ok(objects
            .iter()
            .map(|object| RecognitionObject {
                id: object.id,
                position: object.position,
                orientation: object.orientation,
                size: object.size,
                position_on_image: object.position_on_image,
                size_on_image: object.size_on_image,
                number_of_colors: object.number_of_colors,
                colors: unsafe { from_raw_parts(object.colors, object.number_of_colors as usize) },
                model: unsafe { CStr::from_ptr(object.model) }
                    .to_str()
                    .expect("CStr::to_str"),
            })
            .collect())
    }

    pub fn has_segmentation(&self) -> bool {
        unsafe { wb_camera_recognition_has_segmentation(self.0) != 0 }
    }

    pub fn enable_segmentation(&self) {
        unsafe { wb_camera_recognition_enable_segmentation(self.0) }
    }

    pub fn disable_segmentation(&self) {
        unsafe { wb_camera_recognition_disable_segmentation(self.0) }
    }

    pub fn is_segmentation_enabled(&self) -> bool {
        unsafe { wb_camera_recognition_is_segmentation_enabled(self.0) != 0 }
    }

    pub fn get_segmentation_image(&self) -> &[u8] {
        let camera = Camera::new(self.0);
        let width = camera.get_width();
        let height = camera.get_height();
        unsafe {
            let image = wb_camera_recognition_get_segmentation_image(self.0);
            from_raw_parts(image, (width * height * 4) as usize)
        }
    }

    pub fn save_segmentation_image(&self, filename: &str, quality: i32) -> i32 {
        let filename = CString::new(filename).expect("CString::new failed");
        unsafe { wb_camera_recognition_save_segmentation_image(self.0, filename.as_ptr(), quality) }
    }
}