Skip to content
Snippets Groups Projects
Commit ac7a25af authored by Manon Blanco's avatar Manon Blanco Committed by Yoann Schneider
Browse files

Only predict on folders, no longer support single image

parent b7e3820d
No related branches found
No related tags found
1 merge request!308Only predict on folders, no longer support single image
Hide whitespace changes
Inline Side-by-side
dan/ocr/predict/__init__.py
+ 1
− 7
View file @ ac7a25af
......@@ -17,13 +17,7 @@ def add_predict_parser(subcommands) -> None:
help=__doc__,
)
# Required arguments.
image_or_folder_input = parser.add_mutually_exclusive_group(required=True)
image_or_folder_input.add_argument(
"--image",
type=pathlib.Path,
help="Path to the image to predict.",
)
image_or_folder_input.add_argument(
parser.add_argument(
"--image-dir",
type=pathlib.Path,
help="Path to the folder where the images to predict are stored.",
......
dan/ocr/predict/inference.py
+ 1
− 3
View file @ ac7a25af
......@@ -412,7 +412,6 @@ def process_batch(
def run(
image: Optional[Path],
image_dir: Optional[Path],
model: Path,
output: Path,
......@@ -435,7 +434,6 @@ def run(
) -> None:
"""
Predict a single image save the output
:param image: Path to the image to predict.
:param image_dir: Path to the folder where the images to predict are stored.
:param model: Path to the directory containing the model, the YAML parameters file and the charset file to use for prediction.
:param output: Path to the output folder where the results will be saved.
......@@ -466,7 +464,7 @@ def run(
# Do not use LM with invalid LM weight
use_language_model = dan_model.lm_decoder is not None
images = image_dir.rglob(f"*{image_extension}") if not image else [image]
images = image_dir.rglob(f"*{image_extension}")
for image_batch in list_to_batches(images, n=batch_size):
process_batch(
image_batch,
......
docs/usage/predict/index.md
+ 17
− 18
View file @ ac7a25af
......@@ -6,7 +6,6 @@ Use the `teklia-dan predict` command to apply a trained DAN model on an image.
| Parameter | Description | Type | Default |
| --------------------------- | ------------------------------------------------------------------------------------------------------------------------------------- | -------------- | ------------- |
| `--image` | Path to the image to predict. Must not be provided with `--image-dir`. | `pathlib.Path` | |
| `--image-dir` | Path to the folder where the images to predict are stored. Must not be provided with `--image`. | `pathlib.Path` | |
| `--image-extension` | The extension of the images in the folder. Ignored if `--image-dir` is not provided. | `str` | .jpg |
| `--model` | Path to the directory containing the model, the YAML parameters file and the charset file to use for prediction. | `pathlib.Path` | |
......@@ -41,13 +40,13 @@ To run a prediction with confidence scores, run this command:
```shell
teklia-dan predict \
--image example.jpg \
--image-dir images/ \
--model models \
--output predict/ \
--confidence-score
```
It will create the following JSON file named `predict/example.json`
It will create the following JSON file named after the image in the `predict` folder:
```json
{
......@@ -64,14 +63,14 @@ To run a prediction with confidence scores and plot line-level attention maps, r
```shell
teklia-dan predict \
--image example.jpg \
--image-dir images/ \
--model models \
--output predict/ \
--confidence-score \
--attention-map
```
It will create the following JSON file named `predict/example.json` and a GIF showing a line-level attention map `predict/example_line.gif`
It will create the following JSON file named after the image and a GIF showing a line-level attention map in the `predict` folder:
```json
{
......@@ -91,7 +90,7 @@ To run a prediction with confidence scores and plot word-level attention maps, r
```shell
teklia-dan predict \
--image example.jpg \
--image-dir images/ \
--model models \
--output predict/ \
--confidence-score \
......@@ -100,7 +99,7 @@ teklia-dan predict \
--attention-map-scale 0.5
```
It will create the following JSON file named `predict/example.json` and a GIF showing a word-level attention map `predict/example_word.gif`.
It will create the following JSON file named after the image and a GIF showing a word-level attention map in the `predict` folder:
```json
{
......@@ -120,14 +119,14 @@ To run a prediction, plot line-level attention maps, and extract polygons, run t
```shell
teklia-dan predict \
--image example.jpg \
--image-dir images/ \
--model models \
--output predict/ \
--attention-map \
--predict-objects
```
It will create the following JSON file named `predict/example.json` and a GIF showing a line-level attention map with extracted polygons `predict/example_line.gif`
It will create the following JSON file named after the image and a GIF showing a line-level attention map in the `predict` folder:
```json
{
......@@ -192,13 +191,13 @@ Then, run this command:
```shell
teklia-dan predict \
--image dan_humu_page/6e830f23-e70d-4399-8b94-f36ed3198575.jpg \
--image-dir images/ \
--model models \
--use-language-model \
--output dan_humu_page/predict_char_lm/
--output predict_char_lm/
```
It will create the following JSON file named `dan_humu_page/predict_char_lm/6e830f23-e70d-4399-8b94-f36ed3198575.json`
It will create the following JSON file named after the image in the `predict_char_lm` folder:
```json
{
......@@ -228,13 +227,13 @@ Then, run this command:
```shell
teklia-dan predict \
--image dan_humu_page/6e830f23-e70d-4399-8b94-f36ed3198575.jpg \
--image-dir images/ \
--model models \
--use-language-model \
--output dan_humu_page/predict_subword_lm/
--output predict_subword_lm
```
It will create the following JSON file named `dan_humu_page/predict_subword_lm/6e830f23-e70d-4399-8b94-f36ed3198575.json`
It will create the following JSON file named after the image in the `predict_subword_lm` folder:
```json
{
......@@ -264,13 +263,13 @@ Then, run this command:
```shell
teklia-dan predict \
--image dan_humu_page/6e830f23-e70d-4399-8b94-f36ed3198575.jpg \
--image-dir images/ \
--model models \
--use-language-model \
--output dan_humu_page/predict_word_lm/
--output predict_word_lm/
```
It will create the following JSON file named `dan_humu_page/predict_word_lm/6e830f23-e70d-4399-8b94-f36ed3198575.json`
It will create the following JSON file named after the image in the `predict_word_lm` folder:
```json
{
......
tests/test_prediction.py
+ 10
− 5
View file @ ac7a25af
......@@ -293,9 +293,16 @@ def test_run_prediction(
expected_prediction,
tmp_path,
):
# Make tmpdir and copy needed image inside
image_dir = tmp_path / "images"
image_dir.mkdir()
shutil.copyfile(
(PREDICTION_DATA_PATH / "images" / image_name).with_suffix(".png"),
(image_dir / image_name).with_suffix(".png"),
)
run_prediction(
image=(PREDICTION_DATA_PATH / "images" / image_name).with_suffix(".png"),
image_dir=None,
image_dir=image_dir,
model=PREDICTION_DATA_PATH,
output=tmp_path,
confidence_score=True if confidence_score else False,
......@@ -308,7 +315,7 @@ def test_run_prediction(
temperature=temperature,
predict_objects=False,
max_object_height=None,
image_extension=None,
image_extension=".png",
gpu_device=None,
batch_size=1,
tokens=parse_tokens(PREDICTION_DATA_PATH / "tokens.yml"),
......@@ -488,7 +495,6 @@ def test_run_prediction_batch(
)
run_prediction(
image=None,
image_dir=image_dir,
model=PREDICTION_DATA_PATH,
output=tmp_path,
......@@ -648,7 +654,6 @@ def test_run_prediction_language_model(
yaml.dump(params, (model_path / "parameters.yml").open("w"))
run_prediction(
image=None,
image_dir=image_dir,
model=model_path,
output=tmp_path,
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment