first

.gitignore

@@ -0,0 +1 @@
+/target

Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "yolov8_obb_gui"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+anyhow = "1.0.98"
+image = "0.25"
+usls = { git = "https://github.com/mii443/usls.git", features = [
+    "ort-download-binaries",
+    "directml",
+] }
+windows-capture = "1.4.4"
+eframe = "0.29"
+egui = "0.29"

src/capture.rs

@@ -0,0 +1,129 @@
+use anyhow::Result;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Instant;
+use usls::{Annotator, YOLO};
+use windows_capture::capture::{Context, GraphicsCaptureApiHandler};
+use windows_capture::settings::ColorFormat;
+
+use crate::detection::DetectionResult;
+
+pub struct WindowCapture {
+    model: YOLO,
+    annotator: Annotator,
+    frame_count: u64,
+    detection_result: Arc<Mutex<DetectionResult>>,
+    rgb_buffer: Vec<u8>,
+    shutdown_signal: Arc<AtomicBool>,
+}
+
+impl GraphicsCaptureApiHandler for WindowCapture {
+    type Flags = (
+        YOLO,
+        Annotator,
+        Arc<Mutex<DetectionResult>>,
+        Arc<AtomicBool>,
+    );
+    type Error = Box<dyn std::error::Error + Send + Sync>;
+
+    fn new(ctx: Context<Self::Flags>) -> Result<Self, Self::Error> {
+        let (model, annotator, detection_result, shutdown_signal) = ctx.flags;
+        Ok(Self {
+            model,
+            annotator,
+            frame_count: 0,
+            detection_result,
+            rgb_buffer: Vec::with_capacity(1920 * 1080 * 3),
+            shutdown_signal,
+        })
+    }
+
+    fn on_frame_arrived(
+        &mut self,
+        frame: &mut windows_capture::frame::Frame,
+        capture_control: windows_capture::graphics_capture_api::InternalCaptureControl,
+    ) -> std::result::Result<(), Self::Error> {
+        if self.shutdown_signal.load(Ordering::Relaxed) {
+            println!("Shutdown signal received, stopping capture...");
+            capture_control.stop();
+            return Ok(());
+        }
+
+        self.frame_count += 1;
+
+        let color_format = frame.color_format().clone();
+        let width = frame.width();
+        let height = frame.height();
+        let mut buffer = frame.buffer()?;
+        let raw_data = buffer.as_nopadding_buffer()?;
+
+        let required_size = (width * height * 3) as usize;
+        if self.rgb_buffer.capacity() < required_size {
+            self.rgb_buffer
+                .reserve(required_size - self.rgb_buffer.capacity());
+        }
+        self.rgb_buffer.clear();
+
+        match color_format {
+            ColorFormat::Rgba8 => {
+                self.rgb_buffer.extend(
+                    raw_data
+                        .chunks_exact(4)
+                        .flat_map(|chunk| [chunk[0], chunk[1], chunk[2]]),
+                );
+            }
+            ColorFormat::Bgra8 => {
+                self.rgb_buffer.extend(
+                    raw_data
+                        .chunks_exact(4)
+                        .flat_map(|chunk| [chunk[2], chunk[1], chunk[0]]),
+                );
+            }
+            _ => {
+                eprintln!("Unsupported color format: {:?}", color_format);
+                return Ok(());
+            }
+        };
+
+        let image = usls::Image::from_u8s(&self.rgb_buffer, width, height)?;
+
+        let inference_start = Instant::now();
+        let ys = self.model.forward(&[image.clone()])?;
+        let inference_time = inference_start.elapsed();
+
+        let mut detection_count = 0;
+        let mut confidences = Vec::new();
+
+        for y in &ys {
+            if let Some(obboxes) = y.obbs() {
+                detection_count = obboxes.len();
+                for obbox in obboxes.iter() {
+                    if let Some(conf) = obbox.confidence() {
+                        confidences.push(conf);
+                    }
+                }
+            }
+        }
+
+        let annotated_image = image.clone();
+        let annotated = self.annotator.annotate(&annotated_image, &ys[0])?;
+        let annotated_rgb_data = annotated.to_rgb8().into_raw();
+
+        if let Ok(mut result) = self.detection_result.lock() {
+            result.detection_count = detection_count;
+            result.confidences = confidences;
+            result.inference_time = inference_time;
+            result.annotated_image = Some(Arc::new(annotated_rgb_data));
+            result.image_width = width;
+            result.image_height = height;
+            result.frame_id = self.frame_count;
+        }
+
+        Ok(())
+    }
+
+    fn on_closed(&mut self) -> std::result::Result<(), Self::Error> {
+        println!("Window capture closed, cleaning up resources...");
+        Ok(())
+    }
+}

src/detection.rs

@@ -0,0 +1,26 @@
+use std::sync::Arc;
+
+#[derive(Clone)]
+pub struct DetectionResult {
+    pub detection_count: usize,
+    pub confidences: Vec<f32>,
+    pub inference_time: std::time::Duration,
+    pub annotated_image: Option<Arc<Vec<u8>>>,
+    pub image_width: u32,
+    pub image_height: u32,
+    pub frame_id: u64,
+}
+
+impl Default for DetectionResult {
+    fn default() -> Self {
+        Self {
+            detection_count: 0,
+            confidences: Vec::new(),
+            inference_time: std::time::Duration::from_millis(0),
+            annotated_image: None,
+            image_width: 0,
+            image_height: 0,
+            frame_id: 0,
+        }
+    }
+}

src/gui.rs

@@ -0,0 +1,90 @@
+use eframe::egui;
+use eframe::egui::{ColorImage, ImageData, TextureHandle, TextureOptions};
+use std::sync::{Arc, Mutex};
+use std::time::Instant;
+
+use crate::detection::DetectionResult;
+
+pub struct DetectionApp {
+    detection_result: Arc<Mutex<DetectionResult>>,
+    texture: Option<TextureHandle>,
+    last_update_time: Instant,
+}
+
+impl DetectionApp {
+    pub fn new(detection_result: Arc<Mutex<DetectionResult>>) -> Self {
+        Self {
+            detection_result,
+            texture: None,
+            last_update_time: Instant::now(),
+        }
+    }
+}
+
+impl eframe::App for DetectionApp {
+    fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
+        let now = Instant::now();
+
+        self.last_update_time = now;
+
+        egui::CentralPanel::default().show(ctx, |ui| {
+            ui.heading("YOLOv8 OBB GUI");
+
+            if let Ok(result) = self.detection_result.lock() {
+                ui.separator();
+                ui.label(format!("Detections: {}", result.detection_count));
+                ui.label(format!("Inference Time: {:.2?}", result.inference_time));
+
+                if let Some(ref image_data) = result.annotated_image {
+                    if result.image_width > 0 && result.image_height > 0 {
+                        let color_image = ColorImage::from_rgb(
+                            [result.image_width as usize, result.image_height as usize],
+                            &image_data,
+                        );
+
+                        if let Some(ref mut texture) = self.texture {
+                            texture.set(
+                                ImageData::Color(Arc::new(color_image)),
+                                TextureOptions::default(),
+                            );
+                        } else {
+                            self.texture = Some(ctx.load_texture(
+                                "annotated_image",
+                                ImageData::Color(Arc::new(color_image)),
+                                TextureOptions::default(),
+                            ));
+                        }
+
+                        if let Some(ref texture) = self.texture {
+                            ui.separator();
+                            ui.label("Detection Result:");
+
+                            let max_width = ui.available_width();
+                            let max_height = 600.0;
+                            let scale = (max_width / result.image_width as f32)
+                                .min(max_height / result.image_height as f32)
+                                .min(1.0);
+
+                            let display_width = result.image_width as f32 * scale;
+                            let display_height = result.image_height as f32 * scale;
+
+                            ui.image((texture.id(), egui::vec2(display_width, display_height)));
+                        }
+                    }
+                }
+
+                if !result.confidences.is_empty() {
+                    ui.separator();
+                    ui.label("Confidences:");
+                    for (i, conf) in result.confidences.iter().enumerate() {
+                        ui.label(format!("  Detection {}: {:.2}%", i + 1, conf * 100.0));
+                    }
+                }
+            } else {
+                ui.label("Loading data...");
+            }
+        });
+
+        ctx.request_repaint();
+    }
+}

src/main.rs

@@ -0,0 +1,121 @@
+use anyhow::Result;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::{Arc, Mutex};
+use std::thread;
+use std::time::Instant;
+use usls::{
+    Annotator, Config, Device, SKELETON_COCO_19, SKELETON_COLOR_COCO_19, Scale, Style, Task,
+    Version, YOLO,
+};
+use windows_capture::capture::GraphicsCaptureApiHandler;
+use windows_capture::settings::{ColorFormat, CursorCaptureSettings, DrawBorderSettings, Settings};
+use windows_capture::window::Window;
+
+mod capture;
+mod detection;
+mod gui;
+
+use capture::WindowCapture;
+use detection::DetectionResult;
+use gui::DetectionApp;
+
+fn main() -> Result<()> {
+    let config = Config::yolo()
+        .with_model_file("models/kemomimi.onnx")
+        .with_task(Task::OrientedObjectDetection)
+        .with_version(Version::new(8, 0))
+        .with_scale(Scale::N)
+        .with_model_device(Device::DirectMl(0))
+        .with_class_confs(&[0.2, 0.15])
+        .with_keypoint_confs(&[0.5])
+        .with_topk(5)
+        .with_model_num_dry_run(0);
+
+    println!("Initializing model...");
+    let init_start = Instant::now();
+    let model = YOLO::new(config.commit()?)?;
+    let init_time = init_start.elapsed();
+    println!("Model initialized successfully! (Time: {:.2?})", init_time);
+
+    let annotator = Annotator::default()
+        .with_obb_style(Style::obb().with_draw_fill(false))
+        .with_hbb_style(
+            Style::hbb()
+                .with_draw_fill(false)
+                .with_palette(&usls::Color::palette_coco_80()),
+        )
+        .with_keypoint_style(
+            Style::keypoint()
+                .with_skeleton((SKELETON_COCO_19, SKELETON_COLOR_COCO_19).into())
+                .show_confidence(false)
+                .show_id(true)
+                .show_name(false),
+        )
+        .with_mask_style(Style::mask().with_draw_mask_polygon_largest(true));
+
+    let detection_result = Arc::new(Mutex::new(DetectionResult::default()));
+    let detection_result_clone = detection_result.clone();
+
+    let shutdown_signal = Arc::new(AtomicBool::new(false));
+    let shutdown_signal_clone = shutdown_signal.clone();
+
+    let capture_handle = thread::spawn(move || -> Result<()> {
+        let window = Window::from_name("VRChat")?;
+        let settings = Settings::new(
+            window,
+            CursorCaptureSettings::WithoutCursor,
+            DrawBorderSettings::WithoutBorder,
+            ColorFormat::Rgba8,
+            (
+                model,
+                annotator,
+                detection_result_clone,
+                shutdown_signal_clone,
+            ),
+        );
+
+        println!("Starting window capture...");
+        WindowCapture::start(settings)?;
+        Ok(())
+    });
+
+    let options = eframe::NativeOptions {
+        viewport: eframe::egui::ViewportBuilder::default()
+            .with_inner_size([1200.0, 800.0])
+            .with_title("YOLOv8 OBB GUI"),
+        vsync: false,
+        ..Default::default()
+    };
+
+    println!("Starting GUI application...");
+
+    match eframe::run_native(
+        "YOLOv8 OBB GUI",
+        options,
+        Box::new(|_cc| Ok(Box::new(DetectionApp::new(detection_result)))),
+    ) {
+        Ok(_) => {
+            println!("GUI application closed");
+        }
+        Err(e) => {
+            eprintln!("An error occurred in the GUI application: {}", e);
+        }
+    }
+
+    println!("Application closing, initiating shutdown sequence...");
+
+    shutdown_signal.store(true, Ordering::Relaxed);
+
+    println!("Waiting for capture thread to finish...");
+    match capture_handle.join() {
+        Ok(result) => match result {
+            Ok(_) => println!("Capture thread finished successfully"),
+            Err(e) => eprintln!("Capture thread finished with error: {}", e),
+        },
+        Err(e) => eprintln!("Failed to join capture thread: {:?}", e),
+    }
+
+    println!("Shutdown sequence completed");
+
+    Ok(())
+}