Reorganize code

sgasse · Sep 14, 2024 · 4db2b6b · 4db2b6b
1 parent 8a38407
commit 4db2b6b
Show file tree

Hide file tree

Showing 13 changed files with 47 additions and 489 deletions.
diff --git a/infer_server/src/bin/hour_glass.rs → infer_server/src/bin/infer_server.rs b/infer_server/src/bin/hour_glass.rs → infer_server/src/bin/infer_server.rs
@@ -7,14 +7,12 @@ use axum::{routing::get, Extension, Router};
 use clap::Parser;
 use env_logger::TimestampPrecision;
 use infer_server::{
-    hour_glass::{
-        data_socket::spawn_data_socket,
-        endpoints::{faces_stream, healthcheck, named_stream},
-        inferer::Inferer,
-        router::FrameRouter,
-        INCOMING_FRAMES_CHANNEL, INFER_IMAGES_CHANNEL,
-    },
+    data_socket::spawn_data_socket,
+    endpoints::{faces_stream, healthcheck, named_stream},
+    inferer::Inferer,
     meter::spawn_meter_logger,
+    router::FrameRouter,
+    INCOMING_FRAMES_CHANNEL, INFER_IMAGES_CHANNEL,
 };
 
 #[derive(Parser, Debug)]

diff --git a/infer_server/src/bin/msg_passing.rs b/infer_server/src/bin/msg_passing.rs
diff --git a/infer_server/src/hour_glass/data_socket.rs → infer_server/src/data_socket.rs b/infer_server/src/hour_glass/data_socket.rs → infer_server/src/data_socket.rs
@@ -10,7 +10,7 @@ use tokio::{
 };
 use tokio_util::codec::{Framed, LengthDelimitedCodec};
 
-use crate::hour_glass::StaticFrameSender;
+use crate::StaticFrameSender;
 
 pub async fn spawn_data_socket(
     tx: StaticFrameSender,

diff --git a/infer_server/src/hour_glass/endpoints.rs → infer_server/src/endpoints.rs b/infer_server/src/hour_glass/endpoints.rs → infer_server/src/endpoints.rs
@@ -7,7 +7,7 @@ use futures::StreamExt;
 use serde::Deserialize;
 use tokio_stream::wrappers::BroadcastStream;
 
-use crate::{hour_glass::router::FrameRouter, meter::METER};
+use crate::{meter::METER, router::FrameRouter};
 
 /// Search parameters available to streams.
 #[derive(Debug, Deserialize)]

diff --git a/infer_server/src/hour_glass/mod.rs b/infer_server/src/hour_glass/mod.rs
diff --git a/infer_server/src/hour_glass/inferer.rs → infer_server/src/inferer.rs b/infer_server/src/hour_glass/inferer.rs → infer_server/src/inferer.rs
@@ -7,8 +7,8 @@ use imageproc::{
 use lazy_static::lazy_static;
 
 use crate::{
-    hour_glass::StaticImageReceiver,
     nn::{Bbox, InferModel, UltrafaceModel},
+    StaticImageReceiver,
 };
 
 use super::as_jpeg_stream_item;
@@ -93,7 +93,7 @@ fn draw_bboxes_on_image(
 
 lazy_static! {
     static ref DEJAVU_MONO: rusttype::Font<'static> = {
-        let font_data: &[u8] = include_bytes!("../../../resources/DejaVuSansMono.ttf");
+        let font_data: &[u8] = include_bytes!("../../resources/DejaVuSansMono.ttf");
         let font: rusttype::Font<'static> =
             rusttype::Font::try_from_bytes(font_data).expect("failed to load font");
         font

diff --git a/infer_server/src/lib.rs b/infer_server/src/lib.rs
@@ -6,12 +6,36 @@ use std::{
     hash::{Hash, Hasher},
 };
 
-pub mod hour_glass;
+use bytes::{Bytes, BytesMut};
+use thingbuf::mpsc::{StaticChannel, StaticReceiver, StaticSender};
+
+pub mod data_socket;
+pub mod endpoints;
+pub mod inferer;
 pub mod meter;
-pub mod msg_passing;
 pub mod nn;
+pub mod router;
 pub mod utils;
 
+pub type StaticFrameSender = StaticSender<BytesMut>;
+pub type StaticFrameReceiver = StaticReceiver<BytesMut>;
+
+pub static INCOMING_FRAMES_CHANNEL: StaticChannel<BytesMut, 200> = StaticChannel::new();
+
+pub type BroadcastSender = tokio::sync::broadcast::Sender<Bytes>;
+pub type BroadcastReceiver = tokio::sync::broadcast::Receiver<Bytes>;
+
+pub fn broadcast_channel() -> (BroadcastSender, BroadcastReceiver) {
+    tokio::sync::broadcast::channel(20)
+}
+
+pub type StaticImage = (u32, u32, Vec<u8>, Option<BroadcastSender>);
+
+pub type StaticImageSender = StaticSender<StaticImage>;
+pub type StaticImageReceiver = StaticReceiver<StaticImage>;
+
+pub static INFER_IMAGES_CHANNEL: StaticChannel<StaticImage, 10> = StaticChannel::new();
+
 fn hashed<T>(data: T) -> u64
 where
     T: Hash,
@@ -20,3 +44,14 @@ where
     data.hash(&mut hasher);
     hasher.finish()
 }
+
+fn as_jpeg_stream_item(data: &[u8]) -> Bytes {
+    Bytes::copy_from_slice(
+        &[
+            "--frame\r\nContent-Type: image/jpeg\r\n\r\n".as_bytes(),
+            &data[..],
+            "\r\n\r\n".as_bytes(),
+        ]
+        .concat(),
+    )
+}
diff --git a/infer_server/src/msg_passing/data_socket.rs b/infer_server/src/msg_passing/data_socket.rs
diff --git a/infer_server/src/msg_passing/endpoints.rs b/infer_server/src/msg_passing/endpoints.rs
diff --git a/infer_server/src/msg_passing/mod.rs b/infer_server/src/msg_passing/mod.rs