Merge pull request wormtql#37 from Omico/main

feat: Add --min-level
cdowen · Feb 26, 2022 · da0ffb5 · da0ffb5
2 parents d56be07 + 72b5c1c
commit da0ffb5
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/src/main.rs b/src/main.rs
@@ -82,6 +82,7 @@ fn main() {
         .arg(Arg::with_name("dump").long("dump").required(false).takes_value(false).help("输出模型预测结果、二值化图像和灰度图像，debug专用"))
         .arg(Arg::with_name("capture-only").long("capture-only").required(false).takes_value(false).help("只保存截图，不进行扫描，debug专用"))
         .arg(Arg::with_name("min-star").long("min-star").takes_value(true).help("最小星级"))
+        .arg(Arg::with_name("min-level").long("min-level").takes_value(true).help("最小等级"))
         .arg(Arg::with_name("max-wait-switch-artifact").long("max-wait-switch-artifact").takes_value(true).help("切换圣遗物最大等待时间(ms)"))
         .arg(Arg::with_name("output-dir").long("output-dir").short("o").takes_value(true).help("输出目录").default_value("."))
         .arg(Arg::with_name("scroll-stop").long("scroll-stop").takes_value(true).help("翻页时滚轮停顿时间（ms）（翻页不正确可以考虑加大该选项，默认为80）"))

diff --git a/src/scanner/yas_scanner.rs b/src/scanner/yas_scanner.rs
@@ -22,6 +22,7 @@ pub struct YasScannerConfig {
     max_row: u32,
     capture_only: bool,
     min_star: u32,
+    min_level: u32,
     max_wait_switch_artifact: u32,
     scroll_stop: u32,
     number: u32,
@@ -39,6 +40,7 @@ impl YasScannerConfig {
             capture_only: matches.is_present("capture-only"),
             dump_mode: matches.is_present("dump"),
             min_star: matches.value_of("min-star").unwrap_or("4").parse::<u32>().unwrap(),
+            min_level: matches.value_of("min-level").unwrap_or("0").parse::<u32>().unwrap(),
             max_wait_switch_artifact: matches.value_of("max-wait-switch-artifact").unwrap_or("800").parse::<u32>().unwrap(),
             scroll_stop: matches.value_of("scroll-stop").unwrap_or("80").parse::<u32>().unwrap(),
             number: matches.value_of("number").unwrap_or("0").parse::<u32>().unwrap(),
@@ -433,6 +435,7 @@ impl YasScanner {
         // v bvvmnvbm
         let is_verbose = self.config.verbose;
         let is_dump_mode = self.config.dump_mode;
+        let min_level = self.config.min_level;
         let handle = thread::spawn(move || {
             let mut results: Vec<InternalArtifact> = Vec::new();
             let mut model = CRNNModel::new(
@@ -543,7 +546,14 @@ impl YasScanner {
             info!("error count: {}", error_count);
             info!("dup count: {}", dup_count);
 
-            results
+            if min_level > 0 {
+                results
+                    .into_iter()
+                    .filter(|result| result.level >= min_level)
+                    .collect::<Vec<_>>()
+            } else {
+                results
+            }
         });