From 6430338e952b66a54e8689eb10f08fda3d683f53 Mon Sep 17 00:00:00 2001
From: jhqxxx <jhqxxx@163.com>
Date: Thu, 11 Dec 2025 19:33:30 +0800
Subject: [PATCH] add url crate to handle local path

---
 Cargo.lock             |  1 +
 Cargo.toml             |  3 ++-
 src/utils/img_utils.rs | 18 ++++++++++++++----
 tests/messy_test.rs    | 33 ++++++++++++++++++++++-----------
 4 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 58fab09..d2dd9f3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -43,6 +43,7 @@ dependencies = [
  "serde_json",
  "tokenizers",
  "tokio",
+ "url",
  "uuid",
 ]
 
diff --git a/Cargo.toml b/Cargo.toml
index 5f4f683..17c5420 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -30,6 +30,7 @@ hound = "3.5.1"
 clap = { version = "4.5.51", features = ["derive"] }
 modelscope = "0.1.0"
 dirs = "6.0.0"
+url = "2.5.7"
 
 [features]
 flash-attn=["candle-flash-attn"]
@@ -38,4 +39,4 @@ ffmpeg=["ffmpeg-next"]
 
 [lints.clippy]
 needless_range_loop = "allow"
-single_range_in_vec_init = "allow"
\ No newline at end of file
+single_range_in_vec_init = "allow"
diff --git a/src/utils/img_utils.rs b/src/utils/img_utils.rs
index cf75cc1..99ff264 100644
--- a/src/utils/img_utils.rs
+++ b/src/utils/img_utils.rs
@@ -1,10 +1,10 @@
-use std::collections::HashSet;
 use std::io::Cursor;
+use std::{collections::HashSet, path::PathBuf};
 
 use aha_openai_dive::v1::resources::chat::{
     ChatCompletionParameters, ChatMessage, ChatMessageContent, ChatMessageContentPart,
 };
-use anyhow::{Ok, Result, anyhow};
+use anyhow::{Result, anyhow};
 use base64::{Engine, engine::general_purpose};
 use candle_core::{DType, Device, Tensor};
 use image::{DynamicImage, ImageBuffer, ImageReader, Rgb, RgbImage, imageops};
@@ -48,8 +48,18 @@ pub fn get_image(file: &str) -> Result<DynamicImage> {
         img = Some(load_image_from_url(file)?);
     }
     if file.starts_with("file://") {
-        let mut path = file.to_owned();
-        path = path.split_off(7);
+        // let mut path = file.to_owned();
+        // path = path.split_off(7);
+        let path = url::Url::parse(file)?;
+        let path = path.to_file_path();
+        let path = match path {
+            Ok(path) => path,
+            Err(_) => {
+                let mut path = file.to_owned();
+                path = path.split_off(7);
+                PathBuf::from(path)
+            }
+        };
         img = Some(
             ImageReader::open(path)
                 .map_err(|e| anyhow!(format!("Failed to open file: {}", e)))?
diff --git a/tests/messy_test.rs b/tests/messy_test.rs
index 3b74720..9626571 100644
--- a/tests/messy_test.rs
+++ b/tests/messy_test.rs
@@ -1,20 +1,31 @@
-use std::time::Instant;
+use std::{path::PathBuf, str::FromStr};
 
-use aha::utils::tensor_utils::interpolate_bilinear;
 use anyhow::Result;
-use candle_core::Tensor;
 
 #[test]
 fn messy_test() -> Result<()> {
     // RUST_BACKTRACE=1 cargo test -F cuda messy_test -r -- --nocapture
-    let device = &candle_core::Device::Cpu;
-    let t = Tensor::arange(0.0f32, 40.0, device)?.broadcast_as((1, 1, 40, 40))?;
-    println!("t: {}", t);
-    let i_start = Instant::now();
-    let t_inter = interpolate_bilinear(&t, (20, 20), Some(false))?;
-    let i_duration = i_start.elapsed();
-    println!("Time elapsed in interpolate_bilinear is: {:?}", i_duration);
-    println!("t_inter: {}", t_inter);
+    let path_str = "file://./assets/img/ocr_test1.png";
+    let path = url::Url::from_str(path_str)?;
+    let path = path.to_file_path();
+    let path = match path {
+        Ok(path) => path,
+        Err(_) => {
+            let mut path = path_str.to_owned();
+            path = path.split_off(7);
+            PathBuf::from(path)
+        }
+    };
+    println!("to file path: {:?}", path);
+
+    // let device = &candle_core::Device::Cpu;
+    // let t = Tensor::arange(0.0f32, 40.0, device)?.broadcast_as((1, 1, 40, 40))?;
+    // println!("t: {}", t);
+    // let i_start = Instant::now();
+    // let t_inter = interpolate_bilinear(&t, (20, 20), Some(false))?;
+    // let i_duration = i_start.elapsed();
+    // println!("Time elapsed in interpolate_bilinear is: {:?}", i_duration);
+    // println!("t_inter: {}", t_inter);
     // let x: Vec<u32> = (0..5).flat_map(|_| 0u32..10).collect();
     // let id: Vec<u32> = (0..5).flat_map(|h| vec![h; 10]).collect();
     // println!("x: {:?}", id);