merge pull request #8 from mariinkys/swap_mobile_ocr

OCR Changes: This PR changes the models used for OCR recognition; instead of using the server variants, we use the mobile ones, as the server where we host the application can't handle the load of the server variants. It also makes changes to the OCR processing function by filtering out those results with less than 50% confidence. Finally, I also introduced the ability to remove OCR capabilities by using an environment variable; for now, this does not make the app not load the models; instead, it only disables the OCR button on our frontend. Not loading the models at all would be a good improvement in the future.

Author: mariinkys

.gitignore

@@ -22,4 +22,6 @@ playwright/.cache/
 # ocr_models
 ocr_models/ppocrv5_dict.txt
 ocr_models/ppocrv5_server_det.onnx
-ocr_models/ppocrv5_server_rec.onnx
+ocr_models/ppocrv5_server_rec.onnx
+ocr_models/ppocrv5_mobile_det.onnx
+ocr_models/ppocrv5_mobile_rec.onnx

Dockerfile

@@ -46,15 +46,18 @@ COPY --from=builder /work/dictionaries /app/dictionaries
 
 # Download OCR models and place them under /app/ocr_models
 RUN mkdir -p /app/ocr_models && \
- curl -L -o /app/ocr_models/ppocrv5_server_det.onnx \
- https://github.com/GreatV/oar-ocr/releases/download/v0.1.0/ppocrv5_server_det.onnx && \
- curl -L -o /app/ocr_models/ppocrv5_server_rec.onnx \
- https://github.com/GreatV/oar-ocr/releases/download/v0.1.0/ppocrv5_server_rec.onnx && \
+ curl -L -o /app/ocr_models/ppocrv5_mobile_det.onnx \
+ https://github.com/GreatV/oar-ocr/releases/download/v0.1.0/ppocrv5_mobile_det.onnx && \
+ curl -L -o /app/ocr_models/ppocrv5_mobile_rec.onnx \
+ https://github.com/GreatV/oar-ocr/releases/download/v0.1.0/ppocrv5_mobile_rec.onnx && \
  curl -L -o /app/ocr_models/ppocrv5_dict.txt \
  https://github.com/GreatV/oar-ocr/releases/download/v0.1.0/ppocrv5_dict.txt
 
 ENV RUST_LOG="info"
 ENV LEPTOS_SITE_ADDR="0.0.0.0:8080"
 ENV LEPTOS_SITE_ROOT=./site
+ENV DISABLE_ORC="FALSE"
+
 EXPOSE 8080
+
 CMD ["/app/delphinus"]

ocr_models/README.md

@@ -7,12 +7,12 @@ Models used are:
 ### Text Detection Models
 | Model Type     | Version  | Category | Model File                                                                                                      | Size    | Description                                    |
 |----------------|----------|----------|-----------------------------------------------------------------------------------------------------------------|---------|------------------------------------------------|
-| Text Detection | PP-OCRv5 | Server   | [`ppocrv5_server_det.onnx`](https://github.com/GreatV/oar-ocr/releases/download/v0.1.0/ppocrv5_server_det.onnx) | 87.7MB  | Server variant for high-precision requirements |
+| Text Detection | PP-OCRv5 | Mobile   | [`ppocrv5_mobile_det.onnx`](https://github.com/GreatV/oar-ocr/releases/download/v0.1.0/ppocrv5_mobile_det.onnx) | 4.8MB   | Mobile variant for real-time applications      |
 
 ### Text Recognition Models
 | Model Type       | Version  | Language/Category | Model File                                                                                                              | Size   | Description                      |
 |------------------|----------|-------------------|-------------------------------------------------------------------------------------------------------------------------|--------|----------------------------------|
-| Text Recognition | PP-OCRv5 | Chinese/General   | [`ppocrv5_server_rec.onnx`](https://github.com/GreatV/oar-ocr/releases/download/v0.1.0/ppocrv5_server_rec.onnx)         | 84.1MB | Server variant                   |
+| Text Recognition | PP-OCRv5 | Chinese/General   | [`ppocrv5_mobile_rec.onnx`](https://github.com/GreatV/oar-ocr/releases/download/v0.1.0/ppocrv5_mobile_rec.onnx)         | 16.5MB | Mobile variant                   |
 
 ### Character Dictionaries
 | File Type            | Version        | Category | Model File                                                                                                | Size | Description                  |

src/app.rs

@@ -7,6 +7,7 @@ use leptos_router::{
 
 use crate::{
     components::{NavbarComponent, ToastComponent},
+    config::DelphinusConfig,
     pages::{FaqPage, GeneratorPage, HomePage, NotFound},
 };
 
@@ -15,6 +16,11 @@ pub fn App() -> impl IntoView {
     // Provides context that manages stylesheets, titles, meta tags, etc.
     provide_meta_context();
 
+    // Provides the config as a context
+    let (read, _write) =
+        signal::<DelphinusConfig>(SharedValue::new(DelphinusConfig::get).into_inner());
+    provide_context::<ReadSignal<DelphinusConfig>>(read);
+
     view! {
         // injects a stylesheet into the document <head>
         // id=leptos means cargo-leptos will hot-reload this stylesheet

src/config.rs

@@ -0,0 +1,16 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DelphinusConfig {
+    pub disable_ocr: bool,
+}
+
+impl DelphinusConfig {
+    pub fn get() -> Self {
+        let disable_ocr = std::env::var("DISABLE_OCR")
+            .map(|v| v.to_lowercase() == "true")
+            .unwrap_or(false);
+
+        Self { disable_ocr }
+    }
+}

src/core/flashcard_generation/ocr.rs

@@ -14,12 +14,15 @@ pub async fn ocr_image(bytes: Vec<u8>, separation_char: String) -> Result<String
 
     let mut clean_result = Vec::new();
     for region in &result.text_regions {
-        if let (Some(text), Some(_confidence)) = (&region.text, region.confidence) {
+        if let (Some(text), Some(confidence)) = (&region.text, region.confidence) {
             let filtered = filter_cjk(text);
 
             if !filtered.is_empty() {
-                //println!("Text: {} (confidence: {:.2})", chinese_only, confidence);
-                clean_result.push(filtered);
+                // only those with more than 50% confidence
+                //println!("Text: {} (confidence: {:.2})", filtered, confidence);
+                if confidence > 0.5 {
+                    clean_result.push(filtered);
+                }
             }
         }
     }

src/lib.rs

@@ -1,5 +1,6 @@
 pub mod app;
 pub mod components;
+pub mod config;
 pub mod core;
 pub mod pages;
 

src/main.rs

@@ -1,11 +1,12 @@
 pub mod components;
+pub mod config;
 pub mod core;
 pub mod pages;
 
 #[cfg(feature = "ssr")]
 #[actix_web::main]
 async fn main() -> std::io::Result<()> {
-    use crate::core::flashcard_generation::dictionaries;
+    use crate::{config::DelphinusConfig, core::flashcard_generation::dictionaries};
     use actix_files::Files;
     use actix_web::*;
     use delphinus::app::*;
@@ -31,14 +32,16 @@ async fn main() -> std::io::Result<()> {
 
     let ocr_client = Arc::new(
         OAROCRBuilder::new(
-            String::from("ocr_models/ppocrv5_server_det.onnx"),
-            String::from("ocr_models/ppocrv5_server_rec.onnx"),
+            String::from("ocr_models/ppocrv5_mobile_det.onnx"),
+            String::from("ocr_models/ppocrv5_mobile_rec.onnx"),
             String::from("ocr_models/ppocrv5_dict.txt"),
         )
         .build()
         .expect("Failed to create OCR Client"),
     );
 
+    let _shared_config = SharedValue::new(DelphinusConfig::get);
+
     println!("listening on http://{}", &addr);
 
     HttpServer::new(move || {

src/pages/flashcard_generator.rs

@@ -7,6 +7,7 @@ use crate::{
         DialogComponent, PageTitleComponent, SelectOption, ToastType,
         flashcard_generation::ModifyGeneratedFlashcards, toast::ToastMessage,
     },
+    config::DelphinusConfig,
     core::flashcard_generation::{
         entities::SeparationChar,
         flashcard::{Flashcard, remove_whitespace, search_dictionary},
@@ -16,6 +17,7 @@ use crate::{
 #[component]
 pub fn GeneratorPage() -> impl IntoView {
     let set_toast: WriteSignal<ToastMessage> = expect_context();
+    let delphinus_config: ReadSignal<DelphinusConfig> = expect_context();
 
     let (character_string, set_character_string) = signal(String::new());
     let (language, set_language) = signal("Chinese".to_string());
@@ -189,6 +191,7 @@ pub fn GeneratorPage() -> impl IntoView {
             <div class="max-w-4xl text-center m-auto p-2">
                 <div class="flex justify-between my-3 items-center">
                     <button
+                        disabled=move || {delphinus_config.get().disable_ocr}
                         class="btn btn-sm btn-accent"
                         on:click=move |_| {
                             let _ = ocr_dialog_ref_node.get().unwrap().show_modal();