Skip to content

Commit

Permalink
Add new Argentinian model trained with more (synthetic) data
Browse files Browse the repository at this point in the history
  • Loading branch information
@ankandrew
ankandrew committed May 10, 2024
1 parent 794ff73 commit 6f76560
Show file tree
Hide file tree
Showing 4 changed files with 29 additions and 16 deletions.
7 changes: 4 additions & 3 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,10 @@ The idea is to use this after a plate object detector, since the OCR expects the

### Available Models

| Model Name | Time b=1<br/> (ms)<sup>[1]</sup> | Throughput <br/> (plates/second)<sup>[1]</sup> | Dataset | Accuracy<sup>[2]</sup> | Dataset |
|:----------------------------:|:--------------------------------:|:----------------------------------------------:|:-----------------------------------------------------------------------------------------------------------------------:|:----------------------:|:---------------------------------:|
| argentinian-plates-cnn-model | 2.0964 | 477 | [arg_plate_dataset.zip](https://github.com/ankandrew/fast-plate-ocr/releases/download/arg-plates/arg_plate_dataset.zip) | 94.05% | Non-synthetic, plates up to 2020. |
| Model Name | Time b=1<br/> (ms)<sup>[1]</sup> | Throughput <br/> (plates/second)<sup>[1]</sup> | Dataset | Accuracy<sup>[2]</sup> | Dataset |
|:------------------------------------:|:--------------------------------:|:----------------------------------------------:|:----------------------------------------------------------------------------------------------------------------------------------:|:----------------------:|:-----------------------------------------------:|
| `argentinian-plates-cnn-model` | 2.1 | 476 | [arg_plate_dataset.zip](https://github.com/ankandrew/fast-plate-ocr/releases/download/arg-plates/arg_plate_dataset.zip) | 94.05% | Non-synthetic, plates up to 2020. |
| `argentinian-plates-cnn-synth-model` | 2.1 | 476 | [arg_plate_dataset.zip](https://github.com/ankandrew/fast-plate-ocr/releases/download/arg-plates/arg_plate_dataset_plus_synth.zip) | 94.19% | Plates up to 2020 + Bike & car synthetic plates |

_<sup>[1]</sup> Inference on Mac M1 chip using CPUExecutionProvider. Utilizing CoreMLExecutionProvider accelerates speed
by 5x._
Expand Down
7 changes: 4 additions & 3 deletions docs/index.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,10 @@ The idea is to use this after a plate object detector, since the OCR expects the

We currently have the following available models:

| Model Name | Time b=1<br/> (ms)<sup>[1]</sup> | Throughput <br/> (plates/second)<sup>[1]</sup> | Dataset | Accuracy<sup>[2]</sup> | Dataset |
|:------------------------------:|:--------------------------------:|:----------------------------------------------:|:-----------------------------------------------------------------------------------------------------------------------:|:----------------------:|:---------------------------------:|
| `argentinian-plates-cnn-model` | 2.0964 | 477 | [arg_plate_dataset.zip](https://github.com/ankandrew/fast-plate-ocr/releases/download/arg-plates/arg_plate_dataset.zip) | 94.05% | Non-synthetic, plates up to 2020. |
| Model Name | Time b=1<br/> (ms)<sup>[1]</sup> | Throughput <br/> (plates/second)<sup>[1]</sup> | Dataset | Accuracy<sup>[2]</sup> | Dataset |
|:------------------------------------:|:--------------------------------:|:----------------------------------------------:|:----------------------------------------------------------------------------------------------------------------------------------:|:----------------------:|:-----------------------------------------------:|
| `argentinian-plates-cnn-model` | 2.1 | 476 | [arg_plate_dataset.zip](https://github.com/ankandrew/fast-plate-ocr/releases/download/arg-plates/arg_plate_dataset.zip) | 94.05% | Non-synthetic, plates up to 2020. |
| `argentinian-plates-cnn-synth-model` | 2.1 | 476 | [arg_plate_dataset.zip](https://github.com/ankandrew/fast-plate-ocr/releases/download/arg-plates/arg_plate_dataset_plus_synth.zip) | 94.19% | Plates up to 2020 + Bike & car synthetic plates |

_<sup>[1]</sup> Inference on Mac M1 chip using CPUExecutionProvider. Utilizing CoreMLExecutionProvider accelerates speed
by 5x._
Expand Down
13 changes: 10 additions & 3 deletions fast_plate_ocr/inference/hub.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,18 +7,25 @@
import shutil
import urllib.request
from http import HTTPStatus
from typing import Literal

from tqdm.asyncio import tqdm

from fast_plate_ocr.inference.utils import safe_write

BASE_URL: str = "https://github.com/ankandrew/cnn-ocr-lp/releases/download"
OcrModel = Literal["argentinian-plates-cnn-model", "argentinian-plates-cnn-synth-model"]

AVAILABLE_ONNX_MODELS: dict[str, tuple[str, str]] = {

AVAILABLE_ONNX_MODELS: dict[OcrModel, tuple[str, str]] = {
"argentinian-plates-cnn-model": (
f"{BASE_URL}/arg-plates/arg_cnn_ocr.onnx",
f"{BASE_URL}/arg-plates/arg_cnn_ocr_config.yaml",
)
),
"argentinian-plates-cnn-synth-model": (
f"{BASE_URL}/arg-plates/arg_cnn_ocr_synth.onnx",
f"{BASE_URL}/arg-plates/arg_cnn_ocr_config.yaml",
),
}
"""Available ONNX models for doing inference."""

Expand All @@ -45,7 +52,7 @@ def _download_with_progress(url: str, filename: pathlib.Path) -> None:


def download_model(
model_name: str,
model_name: OcrModel,
save_directory: pathlib.Path | None = None,
force_download: bool = False,
) -> tuple[pathlib.Path, pathlib.Path]:
Expand Down
18 changes: 11 additions & 7 deletions fast_plate_ocr/inference/onnx_inference.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from fast_plate_ocr.common.utils import measure_time
from fast_plate_ocr.inference import hub
from fast_plate_ocr.inference.config import load_config_from_yaml
from fast_plate_ocr.inference.hub import OcrModel
from fast_plate_ocr.inference.process import postprocess_output, preprocess_image, read_plate_image


Expand Down Expand Up @@ -62,11 +63,12 @@ class ONNXPlateRecognizer:

def __init__(
self,
hub_ocr_model: Literal["argentinian-plates-cnn-model"] | None = None,
device: Literal["gpu", "cpu", "auto"] = "auto",
hub_ocr_model: OcrModel | None = None,
device: Literal["cuda", "cpu", "auto"] = "auto",
sess_options: ort.SessionOptions | None = None,
model_path: str | os.PathLike[str] | None = None,
config_path: str | os.PathLike[str] | None = None,
force_download: bool = False,
) -> None:
"""
Initializes the ONNXPlateRecognizer with the specified OCR model and inference device.
Expand All @@ -77,26 +79,26 @@ def __init__(
Args:
hub_ocr_model: Name of the OCR model to use from the HUB.
device: Device type for inference. Should be one of ('cpu', 'gpu', 'auto'). If
device: Device type for inference. Should be one of ('cpu', 'cuda', 'auto'). If
'auto' mode, the device will be deduced from
`onnxruntime.get_available_providers()`.
sess_options: Advanced session options for ONNX Runtime.
model_path: Path to ONNX model file to use (In case you want to use a custom one).
config_path: Path to config file to use (In case you want to use a custom one).
force_download: Force and download the model, even if it already exists.
Returns:
None.
"""
self.logger = logging.getLogger(__name__)

if device == "gpu":
if device == "cuda":
self.provider = ["CUDAExecutionProvider"]
elif device == "cpu":
self.provider = ["CPUExecutionProvider"]
elif device == "auto":
self.provider = ort.get_available_providers()
else:
raise ValueError(f"Device should be one of ('cpu', 'gpu', 'auto'). Got '{device}'.")
raise ValueError(f"Device should be one of ('cpu', 'cuda', 'auto'). Got '{device}'.")

if model_path and config_path:
model_path = pathlib.Path(model_path)
Expand All @@ -106,7 +108,9 @@ def __init__(
self.model_name = model_path.stem
elif hub_ocr_model:
self.model_name = hub_ocr_model
model_path, config_path = hub.download_model(model_name=hub_ocr_model)
model_path, config_path = hub.download_model(
model_name=hub_ocr_model, force_download=force_download
)
else:
raise ValueError(
"Either provide a model from the HUB or a custom model_path and config_path"
Expand Down

0 comments on commit 6f76560

Please sign in to comment.