Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for HDR and Capturing Monitor Output #15

Draft
wants to merge 3 commits into
base: master
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ clap = { version = "4.5.1", features = ["derive"] }
eframe = "0.19.0"
egui_extras = "0.19.0"
global-hotkey = "0.4.2"
image = "0.24.3"
image = { version = "0.24.9", features = ["exr"] }
indexmap = { version = "1.9.1", features = ["serde"] }
lazy_static = "1.4.0"
levenshtein = "1.0.5"
Expand All @@ -35,4 +35,5 @@ serde_json = "1.0.85"
tesseract = "0.12.0"
xcap = "0.0.4"
log = "0.4.22"
env_logger = "0.11.5"
env_logger = "0.11.5"
bytemuck = { version = "1.13", features = ["derive"] }
255 changes: 242 additions & 13 deletions src/bin/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,24 +8,176 @@ use std::{
};
use std::{path::PathBuf, sync::mpsc};

use clap::Parser;
use clap::{Parser};
use env_logger::{Builder, Env};
use global_hotkey::{hotkey::HotKey, GlobalHotKeyEvent, GlobalHotKeyManager, HotKeyState};
use image::DynamicImage;
use log::{debug, error, info, warn};
use notify::{watcher, RecursiveMode, Watcher};
use xcap::Window;
use xcap::{Window, Monitor};
use std::time::{SystemTime, UNIX_EPOCH};
use bytemuck::cast_slice;

use wfinfo::{
database::Database,
ocr::{normalize_string, reward_image_to_reward_names, OCR},
};

fn run_detection(capturer: &Window, db: &Database) {
use image::{DynamicImage, ImageBuffer, ImageFormat, Rgba, RgbaImage, Luma};


fn save_debug_image(image: &DynamicImage, filename: &str) -> Result<(), Box<dyn Error>> {
image.save_with_format(filename, ImageFormat::Png)?;
info!("Saved debug image: {}", filename);
Ok(())
}

fn hdr_to_sdr(image: &ImageBuffer<Rgba<f32>, Vec<f32>>, luminescence: u32) -> ImageBuffer<Rgba<u8>, Vec<u8>> {
let mut sdr_image = ImageBuffer::new(image.width(), image.height());

// Normalize luminescence to 0-1 range
let max_luminescence = 1000.0;
let lum_factor = (luminescence as f32 / max_luminescence).min(1.0).max(0.1);

// Calculate the average luminance of the image
let mut max_luminance = 0.0f32;
for pixel in image.pixels() {
let luminance = 0.2126 * pixel[0] + 0.7152 * pixel[1] + 0.0722 * pixel[2];
max_luminance = max_luminance.max(luminance);
}

// Adjust max_luminance based on the luminescence factor
max_luminance *= lum_factor;

for (x, y, pixel) in image.enumerate_pixels() {
let r = pixel[0];
let g = pixel[1];
let b = pixel[2];

// Apply tone mapping (Reinhard operator)
let luminance = 0.2126 * r + 0.7152 * g + 0.0722 * b;
let scaled_luminance = luminance / max_luminance;
let mapped_luminance = scaled_luminance / (1.0 + scaled_luminance);

// Apply color correction
let scale = mapped_luminance / luminance;
let r_sdr = (r * scale * 255.0).min(255.0) as u8;
let g_sdr = (g * scale * 255.0).min(255.0) as u8;
let b_sdr = (b * scale * 255.0).min(255.0) as u8;

// Apply gamma correction
let gamma = 1.0 / 2.2;
let r_gamma = ((r_sdr as f32 / 255.0).powf(gamma) * 255.0) as u8;
let g_gamma = ((g_sdr as f32 / 255.0).powf(gamma) * 255.0) as u8;
let b_gamma = ((b_sdr as f32 / 255.0).powf(gamma) * 255.0) as u8;

sdr_image.put_pixel(x, y, Rgba([r_gamma, g_gamma, b_gamma, (pixel[3] * 255.0) as u8]));
}

sdr_image
}

fn preprocess_for_ocr(image: &DynamicImage) -> DynamicImage {
let gray_image = image.to_luma8();

// Apply adaptive thresholding
let threshold_image = adaptive_threshold(&gray_image, 11, 2);

DynamicImage::ImageLuma8(threshold_image)
}

fn adaptive_threshold(image: &ImageBuffer<Luma<u8>, Vec<u8>>, block_size: u32, c: i32) -> ImageBuffer<Luma<u8>, Vec<u8>> {
let mut output = ImageBuffer::new(image.width(), image.height());
let half_block = block_size / 2;

for (x, y, pixel) in image.enumerate_pixels() {
let mut sum = 0u32;
let mut count = 0u32;

for i in x.saturating_sub(half_block)..=(x + half_block).min(image.width() - 1) {
for j in y.saturating_sub(half_block)..=(y + half_block).min(image.height() - 1) {
sum += image.get_pixel(i, j).0[0] as u32;
count += 1;
}
}

let threshold = (sum / count) as i32 - c;
let new_value = if pixel.0[0] as i32 > threshold { 255 } else { 0 };
output.put_pixel(x, y, Luma([new_value]));
}

output
}




fn run_detection(capturer: &dyn Capturable, db: &Database, is_hdr: bool, luminescence: u32, save_debug_images: bool) {
let frame = capturer.capture_image().unwrap();
info!("Captured");
let image = DynamicImage::ImageRgba8(frame);
info!("Converted");

let image = if is_hdr {
info!("Converting HDR to SDR");
let converted = DynamicImage::ImageRgba8(hdr_to_sdr(&frame, luminescence));

if save_debug_images {
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();

// Save the original HDR frame
if let Err(e) = image::save_buffer_with_format(
format!("debug_original_{}.exr", timestamp),
cast_slice(frame.as_raw()),
frame.width(),
frame.height(),
image::ColorType::Rgba32F,
image::ImageFormat::OpenExr,
) {
warn!("Failed to save original debug image: {}", e);
}

// Save the converted SDR image
if let Err(e) = save_debug_image(&converted, &format!("debug_converted_{}.png", timestamp)) {
warn!("Failed to save converted debug image: {}", e);
}
}

// Apply preprocessing only for HDR images
let preprocessed = preprocess_for_ocr(&converted);

if save_debug_images {
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
if let Err(e) = save_debug_image(&preprocessed, &format!("debug_preprocessed_{}.png", timestamp)) {
warn!("Failed to save preprocessed debug image: {}", e);
}
}

preprocessed
} else {
DynamicImage::ImageRgba8(rgbaf32_to_rgba8(&frame))
};
info!("Image prepared for OCR");

fn rgbaf32_to_rgba8(rgba: &ImageBuffer<Rgba<f32>, Vec<f32>>) -> ImageBuffer<Rgba<u8>, Vec<u8>> {
let (width, height) = rgba.dimensions();
let mut u8_buffer = Vec::with_capacity((width * height * 4) as usize);

for pixel in rgba.pixels() {
u8_buffer.push((pixel[0] * 255.0) as u8);
u8_buffer.push((pixel[1] * 255.0) as u8);
u8_buffer.push((pixel[2] * 255.0) as u8);
u8_buffer.push((pixel[3] * 255.0) as u8);
}

ImageBuffer::from_raw(width, height, u8_buffer).unwrap()
}



let text = reward_image_to_reward_names(image, None);
let text = text.iter().map(|s| normalize_string(s));
debug!("{:#?}", text);
Expand Down Expand Up @@ -162,13 +314,80 @@ struct Arguments {
/// some systems may require the window name to be specified (e.g. when using gamescope)
#[arg(short, long, default_value = "Warframe")]
window_name: String,
/// Monitor number to capture (0-based index)
///
/// If specified, this will be used instead of the window name
#[arg(short, long)]
monitor: Option<usize>,
/// Specify if the monitor is in HDR mode
#[arg(long)]
hdr: bool,
/// Luminescence level for HDR (100-1000 nits)
#[arg(long, default_value = "300")]
luminescence: u32,
/// Save debug images when HDR conversion is applied
#[arg(long)]
save_debug_images: bool,
}

trait Capturable {
fn capture_image(&self) -> Result<image::ImageBuffer<Rgba<f32>, Vec<f32>>, Box<dyn Error>>;
fn width(&self) -> u32;
fn height(&self) -> u32;
}


impl Capturable for Window {
fn capture_image(&self) -> Result<ImageBuffer<Rgba<f32>, Vec<f32>>, Box<dyn Error>> {
let rgba_image: RgbaImage = self.capture_image()?;
let float_image = rgba_to_rgbaf32(&rgba_image);
Ok(float_image)
}

fn width(&self) -> u32 {
self.width()
}

fn height(&self) -> u32 {
self.height()
}
}

impl Capturable for Monitor {
fn capture_image(&self) -> Result<ImageBuffer<Rgba<f32>, Vec<f32>>, Box<dyn Error>> {
let rgba_image: RgbaImage = self.capture_image()?;
let float_image = rgba_to_rgbaf32(&rgba_image);
Ok(float_image)
}

fn width(&self) -> u32 {
self.width()
}

fn height(&self) -> u32 {
self.height()
}
}

fn rgba_to_rgbaf32(rgba: &RgbaImage) -> ImageBuffer<Rgba<f32>, Vec<f32>> {
let (width, height) = rgba.dimensions();
let mut float_buffer = Vec::with_capacity((width * height * 4) as usize);

for pixel in rgba.pixels() {
float_buffer.push(pixel[0] as f32 / 255.0);
float_buffer.push(pixel[1] as f32 / 255.0);
float_buffer.push(pixel[2] as f32 / 255.0);
float_buffer.push(pixel[3] as f32 / 255.0);
}

ImageBuffer::from_raw(width, height, float_buffer).unwrap()
}


fn main() -> Result<(), Box<dyn Error>> {
let arguments = Arguments::parse();
let default_log_path = PathBuf::from_str(&std::env::var("HOME").unwrap()).unwrap().join(PathBuf::from_str(".local/share/Steam/steamapps/compatdata/230410/pfx/drive_c/users/steamuser/AppData/Local/Warframe/EE.log")?);
let log_path = arguments.game_log_file_path.unwrap_or(default_log_path);
let window_name = arguments.window_name;
let env = Env::default()
.filter_or("WFINFO_LOG", "info")
.write_style_or("WFINFO_STYLE", "always");
Expand All @@ -179,16 +398,26 @@ fn main() -> Result<(), Box<dyn Error>> {
.format_target(false)
.init();

let windows = Window::all()?;
let db = Database::load_from_file(None, None);
let Some(warframe_window) = windows.iter().find(|x| x.title() == window_name) else {
return Err("Warframe window not found".into());

let capturer: Box<dyn Capturable> = if let Some(monitor_index) = arguments.monitor {
let monitors = Monitor::all()?;
if monitor_index >= monitors.len() {
return Err(format!("Invalid monitor index: {}", monitor_index).into());
}
Box::new(monitors[monitor_index].clone())
} else {
let windows = Window::all()?;
let Some(warframe_window) = windows.iter().find(|x| x.title() == arguments.window_name) else {
return Err("Warframe window not found".into());
};
Box::new(warframe_window.clone())
};

debug!(
"Capture source resolution: {:?}x{:?}",
warframe_window.width(),
warframe_window.height()
capturer.width(),
capturer.height()
);

info!("Loaded database");
Expand All @@ -200,7 +429,7 @@ fn main() -> Result<(), Box<dyn Error>> {

while let Ok(()) = event_receiver.recv() {
info!("Capturing");
run_detection(warframe_window, &db);
run_detection(&*capturer, &db, arguments.hdr, arguments.luminescence, arguments.save_debug_images);
}

drop(OCR.lock().unwrap().take());
Expand Down
Loading