[C] Update README and infer module

README.md

@@ -24,7 +24,151 @@ On our validation dataset, our model demonstrated over 99% accuracy based on a z
 
 ---
 
-## Model Evaluation (Benchmark)
+## Table of Contents
+
+- [Introduction](#introduction)
+- [Table of Contents](#table-of-contents)
+- [Quick Start](#quick-start)
+- [Benchmark](#benchmark)
+- [Before We start Training](#before-we-start-training)
+- [Training the Model](#training-the-model)
+- [Model Architecture Design](#model-architecture-design)
+- [Dataset](#dataset)
+- [Dataset Implementation](#dataset-implementation)
+- [Building the Training Environment](#building-the-training-environment)
+- [Running Training (Based on Docker)](#running-training-based-on-docker)
+- [Convert to ONNX Format](#convert-to-onnx-format)
+- [Dataset Submission](#dataset-submission)
+- [Frequently Asked Questions (FAQs)](#frequently-asked-questions-faqs)
+- [Citation](#citation)
+
+---
+
+## Quick Start
+
+### Installation
+
+Currently, we do not provide an installation package on Pypi. If you want to use this project, you can clone it directly from GitHub and then install the required dependencies. Before installing, please ensure that you have already installed [DocsaidKit](https://github.com/DocsaidLab/DocsaidKit).
+
+If you have already installed DocsaidKit, please follow the steps below:
+
+1. Clone the Project:
+
+   ```bash
+   git clone https://github.com/DocsaidLab/DocClassifier.git
+   ```
+
+2. Enter the Project Directory:
+
+   ```bash
+   cd DocClassifier
+   ```
+
+3. Create Packaging Files:
+
+   ```bash
+   python setup.py bdist_wheel
+   ```
+
+4. Install the Packaged Files:
+
+   ```bash
+   pip install dist/docclassifier-*-py3-none-any.whl
+   ```
+
+By following these steps, you should be able to successfully complete the installation of DocClassifier.
+
+Once the installation is complete, you can start using the project.
+
+---
+
+### Importing Necessary Dependencies
+
+We provide a simple model inference interface that includes the logic for preprocessing and postprocessing.
+
+First, you need to import the required dependencies and create the DocClassifier class.
+
+```python
+import docsaidkit as D
+from docsaidkit import Backend
+from docclassifier import DocClassifier
+```
+
+### DocBank
+
+In the inference folder directory, there is a `docbank` folder containing all the registered data. You can place your registration data in it. During inference, specifying the `docbank`, DocClassifier will automatically read all the data in the folder.
+
+If you want to use your own dataset, when creating DocClassifier, please specify the `docbank_root` parameter and set it as your dataset root directory. We recommend that your data uses full-page images and minimizes background interference to enhance the stability of the model.
+
+### ModelType
+
+`ModelType` is an enumeration type used to specify the model type used by DocClassifier. It includes the following options:
+
+- `margin_based`: Uses a model architecture based on the margin method.
+
+There may be more model types in the future, and we will update here.
+
+### Backend
+
+`Backend` is an enumeration type used to specify the computational backend of DocClassifier. It includes the following options:
+
+- `cpu`: Uses CPU for computation.
+- `cuda`: Uses GPU for computation (requires appropriate hardware support).
+
+ONNXRuntime supports a wide range of backends, including CPU, CUDA, OpenCL, DirectX, TensorRT, etc. If you have other requirements, you can refer to [**ONNXRuntime Execution Providers**](https://onnxruntime.ai/docs/execution-providers/index.html) and modify it to the corresponding backend accordingly.
+
+### Creating a DocClassifier Instance
+
+```python
+model = DocClassifier(
+    gpu_id=0,  # GPU ID, set to -1 if not using GPU
+    backend=Backend.cpu,  # Choose the computation backend, can be Backend.cpu or Backend.cuda
+    threshold=0.5,  # Threshold for model prediction, each model has a default value, no need to set if not adjusting
+    docbank_root='path/to/your/docbank',  # Root directory for the registration data, default is docbank
+)
+```
+
+Notes:
+
+- Using CUDA for computation requires not only appropriate hardware support but also the installation of the corresponding CUDA drivers and CUDA Toolkit. If your system does not have CUDA installed, or if the installed version is incorrect, the CUDA computation backend cannot be used.
+
+- For issues related to ONNXRuntime installation dependencies, please refer to [ONNXRuntime Release Notes](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements).
+
+### Reading and Processing Images
+
+```python
+# Reading the image
+img = D.imread('path/to/your/image.jpg')
+
+# Using the model for inference
+# Here, most_similar is the most similar category, max_score is the highest score
+most_similar, max_score = model(img)
+```
+
+### Output Results
+
+The inference result you get is based on the registered dataset within the module. The model will find the most similar category from it and output the label and score of the category.
+
+```python
+import docsaidkit as D
+from docclassifier import DocClassifier
+
+model = DocClassifier(backend=D.Backend.cpu)
+img = D.imread('docs/test_driver.jpg')
+most_similar, max_score = model(img)
+print(f'most_similar: {most_similar}, max_score: {max_score:.4f}')
+# >>> most_similar: Taiwanese Driver's License Front, max_score: 0.7334
+```
+
+- **most_similar: Taiwanese Driver's License Front, max_score: 0.7334**
+
+    <div>
+        <img src="./docs/test_driver.jpg" width="300">
+    </div>
+
+---
+
+## Benchmark
 
 We have an internal test dataset, but due to privacy protection, we cannot make this dataset open source. We can only provide evaluation results based on this dataset.
 
@@ -85,7 +229,7 @@ We have an internal test dataset, but due to privacy protection, we cannot make
     </div>
 ---
 
-## Before Starting Model Training
+## Before We start Training
 
 Based on the model we provide, we believe it can solve most application scenarios. However, we also understand that some scenarios may require enhanced model performance, necessitating the collection of specific datasets and model fine-tuning. We empathize that you may have the budget but not the time to customize adjustments to your specific environment. Therefore, you can contact us directly for consultation. Depending on the complexity of your project, we can arrange for engineers to develop custom solutions.
 
@@ -173,7 +317,7 @@ Reference: [ArcFace: Additive Angular Margin Loss for Deep Face Recognition](htt
 
 ---
 
-## Dataset Introduction and Preprocessing
+## Dataset
 
 Most of the text images are sourced from internet searches.
 
@@ -363,7 +507,7 @@ Here is our default [Dockerfile](./docker/Dockerfile), specifically designed for
 
 ---
 
-## Executing Training (Based on Docker)
+## Running Training (Based on Docker)
 
 Here's how you can use your pre-built Docker image to run document classification training.
 
@@ -427,7 +571,7 @@ Let me know if you need any more assistance!
 
 ---
 
-## Converting the Model to ONNX Format (Based on Docker)
+## Convert to ONNX Format
 
 This section describes how to convert your model to the ONNX format.
 

README_cn.md

@@ -26,6 +26,22 @@ DocClassifier 是一個創新的文件圖像分類系統，它的設計靈感源
 
 ## 目錄
 
+- [介紹](#介紹)
+- [目錄](#目錄)
+- [快速開始](#快速開始)
+- [評估模型（Benchmark）](#評估模型benchmark)
+- [開始訓練模型之前](#開始訓練模型之前)
+- [訓練模型](#訓練模型)
+- [模型架構設計](#模型架構設計)
+- [資料集介紹及預處理](#資料集介紹及預處理)
+- [資料集實作](#資料集實作)
+- [構建訓練環境](#構建訓練環境)
+- [執行訓練（Based on Docker）](#執行訓練based-on-docker)
+- [轉換模型為 ONNX 格式（Based on Docker）](#轉換模型為-onnx-格式based-on-docker)
+- [提交資料集](#提交資料集)
+- [常見問題（FAQs）](#常見問題faqs)
+- [引用](#引用)
+
 ---
 
 ## 快速開始
@@ -79,6 +95,20 @@ from docsaidkit import Backend
 from docclassifier import DocClassifier
 ```
 
+### DocBank
+
+在推論的資料夾目錄中，有一個 `docbank` 資料夾，裡面包含了所有的註冊資料，您可以在其中放置您的註冊資料，並且在推論時，指定 `docbank` DocClassifier 會自動讀取資料夾中的所有資料。
+
+如果您要使用自己的資料集，在創建 DocClassifier 時，請指定 `docbank_root` 參數，並且將其設定為您的資料集根目錄。我們建議您的資料使用滿版的圖像，盡量減少背景的干擾，以提高模型的穩定性。
+
+### ModelType
+
+`ModelType` 是一個枚舉類型，用於指定 DocClassifier 使用的模型類型。它包含以下選項：
+
+- `margin_based`：使用基於 margin 方法的模型架構。
+
+未來可能會有更多的模型類型，我們會在此處更新。
+
 ### Backend
 
 `Backend` 是一個枚舉類型，用於指定 DocClassifier 的運算後端。它包含以下選項：
@@ -94,6 +124,8 @@ ONNXRuntime 支援了非常多的後端，包括 CPU、CUDA、OpenCL、DirectX
 model = DocClassifier(
     gpu_id=0,  # GPU 編號，如果不使用 GPU 請設為 -1
     backend=Backend.cpu,  # 選擇運算後端，可以是 Backend.cpu 或 Backend.cuda
+    threshold=0.5,  # 模型預測的閾值，每個模型都有預設值，如果不需要調整，可以不用設定
+    docbank_root='path/to/your/docbank',  # 註冊資料的根目錄，預設為 docbank
 )
 ```
 
@@ -109,11 +141,9 @@ model = DocClassifier(
 # 讀取圖像
 img = D.imread('path/to/your/image.jpg')
 
-# 您也可以使用我們提供的測試圖像
-# img = D.imread('docs/run_test_card.jpg')
-
 # 使用模型進行推論
-result = model(img) # result 是一個 Document 類型
+# 其中，most_similar 是最相似的類別，max_score 是最相似的分數
+most_similar, max_score = model(img)
 ```
 
 ### 輸出結果
@@ -124,14 +154,18 @@ result = model(img) # result 是一個 Document 類型
 import docsaidkit as D
 from docclassifier import DocClassifier
 
-model = DocClassifier(D.Backend.cpu)
-img = D.imread('docs/run_test_card.jpg')
-result = model(img)
+model = DocClassifier(backend=D.Backend.cpu)
+img = D.imread('docs/test_driver.jpg')
+most_similar, max_score = model(img)
+print(f'most_similar: {most_similar}, max_score: {max_score:.4f}')
+# >>> most_similar: 台灣駕照正面, max_score: 0.7334
 ```
 
-<div align="center">
-    <img src="./docs/result.jpg" width="800">
-</div>
+- **most_similar: 台灣駕照正面, max_score: 0.7334**
+
+    <div>
+        <img src="./docs/test_driver.jpg" width="300">
+    </div>
 
 ---
 

docclassifier/__init__.py

@@ -1 +1,3 @@
 from .classifier import DocClassifier
+
+__version__ = '0.1.0'

docclassifier/classifier.py

@@ -23,6 +23,8 @@ def __init__(
         model_cfg: str = None,
         backend: D.Backend = D.Backend.cpu,
         gpu_id: int = 0,
+        threshold: float = None,
+        docbank_root: Union[str, D.Path] = None,
         **kwargs
     ):
         model_type = ModelType.obj_to_enum(model_type)
@@ -38,6 +40,8 @@ def __init__(
                 gpu_id=gpu_id,
                 backend=backend,
                 model_cfg=model_cfg,
+                threshold=threshold,
+                docbank_root=docbank_root,
                 **kwargs
             )
 

docclassifier/margin_based/infer.py

@@ -38,11 +38,13 @@ class Inference:
             'model_path': 'lcnet050_arcface_fp32.onnx',
             'file_id': '',
             'img_size_infer': (96, 96),
+            'threshold': 0.7021  # FPR=0.01
         },
         'lcnet050_cosface': {
             'model_path': 'lcnet050_cosface_fp32.onnx',
             'file_id': '',
             'img_size_infer': (96, 96),
+            'threshold': 0.6765  # FPR=0.01
         },
     }
 
@@ -51,7 +53,7 @@ def __init__(
         gpu_id: int = 0,
         backend: D.Backend = D.Backend.cpu,
         model_cfg: str = 'lcnet050_arcface',
-        threshold: float = 0.8,
+        threshold: float = None,
         docbank_root: Union[str, D.Path] = None,
         **kwargs
     ):
@@ -64,7 +66,7 @@ def __init__(
             os.system(D.gen_download_cmd(file_id, model_path))
         self.model = D.ONNXEngine(model_path, gpu_id, backend, **kwargs)
         self.bank = self.get_doc_bank(docbank_root)
-        self.threshold = threshold
+        self.threshold = threshold if threshold is not None else cfg['threshold']
 
     @staticmethod
     def compare(feat1, feat2):