core/

sauron.rs

//! Sauron is an actor which implements the object detection and localization pipeline.

use crate::database::{create_schema, get_gps, smart_update_database};
use crate::mailbox::MailboxHandle;
use crate::utils::get_env;

use chrono::{DateTime, Utc};
use image::ImageReader;
use serde_json::json;
use sqlx::{PgPool, postgres::PgPoolOptions};
use std::path::Path;
use tokio::sync::mpsc;
use tracing::info;

use algorithms::localize;
use sauron::yolo;
use types::{
    cv::{Gps, Model, Object},
    mailbox::{MailboxMessage, MailboxMessageType},
};

/// Sauron is an actor which implements the object detection and localization pipeline.
struct Sauron {
    receiver: mpsc::Receiver<SauronMessage>,
    model: Model,
    pool: PgPool,
    mailbox: Option<MailboxHandle>,
}

#[allow(dead_code)]
enum SauronMessage {
    RunSauron { file_path: String, timestamp: i64 },
    SetMailbox { mailbox: MailboxHandle },
}

impl Sauron {
    /// Create a new instance of Sauron.
    ///
    /// To create a Sauron actor, call `SauronHandle::new()`.
    async fn new(receiver: mpsc::Receiver<SauronMessage>) -> Self {
        info!("Sauron Started");
        let database_url = get_env(
            "DATABASE_URL",
            "postgresql://user:password@localhost:5432/local".to_string(),
        );
        info!("Sauron Connecting with url: {database_url}");

        let pool = PgPoolOptions::new()
            .max_connections(5)
            .connect(&database_url)
            .await
            .expect("Failed to create pool.");

        if let Err(e) = create_schema(&pool).await {
            eprintln!("Failed to create database schema: {}", e);
        }

        let model = yolo::load_model().expect("Loading Model Failed");

        Sauron {
            receiver,
            model,
            pool,
            mailbox: None,
        }
    }

    /// Process a message received by the Sauron actor.
    async fn handle_message(&mut self, msg: SauronMessage) {
        match msg {
            SauronMessage::RunSauron {
                file_path,
                timestamp,
            } => {
                self.run_odlc(file_path, timestamp).await.unwrap_or(());
            }
            SauronMessage::SetMailbox { mailbox } => {
                self.mailbox = Some(mailbox);
                info!("Mailbox connected to Sauron");
            }
        }
    }

    /// Runs the object detection and localization pipeline.
    ///
    /// # Arguments
    ///
    /// * `file_path` - The path to the image file to localize
    /// * `timestamp` - The timestamp of the image
    async fn run_odlc(&self, file_path: String, timestamp: i64) -> anyhow::Result<()> {
        //Runs Odlc pipeline
        info!("Running ODLC pipeline for {}", file_path);

        // Run detect
        let image = ImageReader::open(Path::new(&file_path))?.decode()?;
        let detections = yolo::detect(&self.model, &image).await?;

        // get gps and interpolate them
        let (gps1, gps2) = get_gps(&self.pool, &timestamp).await?;
        let gps = Sauron::interpolate_gps(
            &gps1.into(),
            &gps2.into(),
            DateTime::from_timestamp_millis(timestamp)
                .expect("Failed to convert image timestamp to datetime"),
        )
        .await;

        // localize each detection
        let objects: Vec<Object> = detections
            .iter()
            .map(|detection| localize(gps, detection))
            .collect();

        // update objects in database
        smart_update_database(&self.pool, objects.clone(), &file_path).await?;

        // Send target found messages to mailbox
        if let Some(ref mailbox) = self.mailbox {
            for object in objects {
                let message = MailboxMessage {
                    msg_type: MailboxMessageType::TargetFound,
                    data: json!({
                        "lat": object.lat,
                        "long": object.long,
                        "class": object.class,
                        "confidence": object.confidence,
                        "image": file_path,
                        "timestamp": timestamp,
                    }),
                };
                mailbox.queue_message(message).await;
            }
        }

        Ok(())
    }

    /// Returns GPS interpolating from the two given GPS points.
    pub async fn interpolate_gps(a: &Gps, b: &Gps, time: DateTime<Utc>) -> Gps {
        let difference = time - a.time;
        let total = b.time - a.time;
        let scale = (difference.num_milliseconds() as f64) / (total.num_milliseconds() as f64);

        let lat = a.lat + (b.lat - a.lat) * scale;
        let long = a.lat + (b.lat - a.lat) * scale;
        let alt = a.alt + (b.alt - a.alt) * scale;
        let heading = a.heading + (b.heading - a.heading) * scale;

        Gps {
            lat,
            long,
            alt,
            heading,
            time,
        }
    }
}

/// Sauron's handle. Used to interact with the actor.
#[derive(Clone)]
pub struct SauronHandle {
    sender: mpsc::Sender<SauronMessage>,
}

impl SauronHandle {
    /// Creates a new Sauron actor, and returns its handle.
    pub async fn new() -> Self {
        let (sender, receiver) = mpsc::channel(8);
        let actor = Sauron::new(receiver).await;
        tokio::spawn(run_sauron(actor));

        Self { sender }
    }

    /// Sends a photo to Sauron for processing.
    pub async fn send_photo(&self, file_path: String, timestamp: i64) {
        let msg = SauronMessage::RunSauron {
            file_path,
            timestamp,
        };

        let _ = self.sender.send(msg).await;
    }

    /// Sets the handle for the mailbox actor which Sauron will send telemetry to.
    #[allow(dead_code)]
    pub async fn set_mailbox(&self, mailbox: MailboxHandle) {
        let msg = SauronMessage::SetMailbox { mailbox };
        let _ = self.sender.send(msg).await;
    }
}

/// Internal function which runs the Sauron actor.
async fn run_sauron(mut actor: Sauron) {
    while let Some(msg) = actor.receiver.recv().await {
        actor.handle_message(msg).await;
    }
}