diff --git a/dan/manager/dataset.py b/dan/manager/dataset.py
index 4afb776c9389ed7c9c7e6003a1a7e3068d5187ce..7317043ffc8c0251b58d5a66309963a4bdb27498 100644
--- a/dan/manager/dataset.py
+++ b/dan/manager/dataset.py
@@ -4,12 +4,10 @@ import json
import os
import numpy as np
-import torch
from torch.utils.data import Dataset
-from torchvision.io import ImageReadMode, read_image
from dan.datasets.utils import natural_sort
-from dan.utils import token_to_ind
+from dan.utils import read_image, token_to_ind
class OCRDataset(Dataset):
@@ -86,14 +84,6 @@ class OCRDataset(Dataset):
return sample
- @staticmethod
- def load_image(path):
- """
- Load an image as a torch.Tensor and scale the values between 0 and 1.
- """
- img = read_image(path, mode=ImageReadMode.RGB)
- return img.to(dtype=torch.get_default_dtype()).div(255)
-
def load_samples(self, paths_and_sets):
"""
Load images and labels
@@ -120,7 +110,7 @@ class OCRDataset(Dataset):
)
if self.load_in_memory:
samples[-1]["img"] = self.preprocessing_transforms(
- self.load_image(filename)
+ read_image(filename)
)
return samples
@@ -130,10 +120,8 @@ class OCRDataset(Dataset):
"""
if self.load_in_memory:
return self.samples[i]["img"]
- else:
- return self.preprocessing_transforms(
- self.load_image(self.samples[i]["path"])
- )
+
+ return self.preprocessing_transforms(read_image(self.samples[i]["path"]))
def compute_std_mean(self):
"""
diff --git a/dan/manager/ocr.py b/dan/manager/ocr.py
index 8391eb7888ea3b23604d4d5dab815194e71c1b79..077f965cb5d9d1ec35300294b033cea14dd9d51b 100644
--- a/dan/manager/ocr.py
+++ b/dan/manager/ocr.py
@@ -62,7 +62,7 @@ class OCRDatasetManager:
else None
)
self.preprocessing = get_preprocessing_transforms(
- params["config"]["preprocessings"]
+ params["config"]["preprocessings"], to_pil_image=True
)
def load_datasets(self):
diff --git a/dan/predict/__init__.py b/dan/predict/__init__.py
index fa81b61e8039dfa13c833ffc7f59d0f5d498b541..1b73e3b8ec20ca7ff87d42947e1469933683cdcb 100644
--- a/dan/predict/__init__.py
+++ b/dan/predict/__init__.py
@@ -58,20 +58,6 @@ def add_predict_parser(subcommands) -> None:
help="The extension of the images in the folder.",
default=".jpg",
)
- parser.add_argument(
- "--scale",
- type=float,
- default=1.0,
- required=False,
- help="Image scaling factor before feeding it to DAN",
- )
- parser.add_argument(
- "--image-max-width",
- type=int,
- default=None,
- required=False,
- help="Image resizing before feeding it to DAN",
- )
parser.add_argument(
"--temperature",
type=float,
diff --git a/dan/predict/prediction.py b/dan/predict/prediction.py
index 0db9b6c987929a19b50ec24bc42ce9beeb322e41..1c1d02afb4a0820a4a8291319fb6331c694c04d8 100644
--- a/dan/predict/prediction.py
+++ b/dan/predict/prediction.py
@@ -5,7 +5,6 @@ import pickle
from itertools import pairwise
from pathlib import Path
-import cv2
import numpy as np
import torch
import yaml
@@ -20,6 +19,7 @@ from dan.predict.attention import (
plot_attention,
split_text_and_confidences,
)
+from dan.transforms import get_preprocessing_transforms
from dan.utils import ind_to_token, read_image
@@ -74,23 +74,27 @@ class DAN:
self.encoder = encoder
self.decoder = decoder
- self.mean, self.std = parameters["mean"], parameters["std"]
+ self.mean, self.std = (
+ torch.tensor(parameters["mean"]) / 255,
+ torch.tensor(parameters["std"]) / 255,
+ )
+ self.preprocessing_transforms = get_preprocessing_transforms(
+ parameters.get("preprocessings", [])
+ )
self.max_chars = parameters["max_char_prediction"]
- def preprocess(self, input_image):
+ def preprocess(self, path):
"""
- Preprocess an input_image.
- :param input_image: The input image to preprocess.
+ Preprocess an image.
+ :param path: Path of the image to load and preprocess.
"""
- assert isinstance(
- input_image, np.ndarray
- ), "Input image must be an np.array in RGB"
- input_image = np.asarray(input_image)
- if len(input_image.shape) < 3:
- input_image = cv2.cvtColor(input_image, cv2.COLOR_GRAY2RGB)
-
- input_image = (input_image - self.mean) / self.std
- return input_image
+ image = read_image(path)
+ preprocessed_image = self.preprocessing_transforms(image)
+ for ch in range(preprocessed_image.shape[0]):
+ preprocessed_image[ch, :, :] = (
+ preprocessed_image[ch, :, :] - self.mean[ch]
+ ) / self.std[ch]
+ return preprocessed_image
def predict(
self,
@@ -252,11 +256,10 @@ class DAN:
def process_image(
- image,
+ image_path,
dan_model,
device,
output,
- scale,
confidence_score,
confidence_score_levels,
attention_map,
@@ -264,27 +267,18 @@ def process_image(
attention_map_scale,
word_separators,
line_separators,
- image_max_width,
predict_objects,
threshold_method,
threshold_value,
):
# Load image and pre-process it
- if image_max_width:
- _, w, _ = read_image(image, scale=1).shape
- ratio = image_max_width / w
- im = read_image(image, ratio)
- else:
- im = read_image(image, scale=scale)
-
+ image = dan_model.preprocess(str(image_path))
logger.info("Image loaded.")
- im_p = dan_model.preprocess(im)
- logger.debug("Image pre-processed.")
# Convert to tensor of size (batch_size, channel, height, width) with batch_size=1
- input_tensor = torch.from_numpy(im_p).permute(2, 0, 1).unsqueeze(0)
+ input_tensor = image.unsqueeze(0)
input_tensor = input_tensor.to(device)
- input_sizes = [im_p.shape[:2]]
+ input_sizes = [image.shape[1:]]
# Parse delimiters to regex
word_separators = parse_delimiters(word_separators)
@@ -346,11 +340,11 @@ def process_image(
# Save gif with attention map
if attention_map:
- gif_filename = f"{output}/{image.stem}_{attention_map_level}.gif"
+ gif_filename = f"{output}/{image_path.stem}_{attention_map_level}.gif"
logger.info(f"Creating attention GIF in {gif_filename}")
# this returns polygons but unused for now.
plot_attention(
- image=im,
+ image=image,
text=prediction["text"][0],
weights=prediction["attentions"][0],
level=attention_map_level,
@@ -364,7 +358,7 @@ def process_image(
)
result["attention_gif"] = gif_filename
- json_filename = f"{output}/{image.stem}.json"
+ json_filename = f"{output}/{image_path.stem}.json"
logger.info(f"Saving JSON prediction in {json_filename}")
save_json(Path(json_filename), result)
@@ -376,7 +370,6 @@ def run(
parameters,
charset,
output,
- scale,
confidence_score,
confidence_score_levels,
attention_map,
@@ -385,7 +378,6 @@ def run(
word_separators,
line_separators,
temperature,
- image_max_width,
predict_objects,
threshold_method,
threshold_value,
@@ -400,14 +392,12 @@ def run(
:param parameters: Path to the YAML parameters file.
:param charset: Path to the charset.
:param output: Path to the output folder where the results will be saved.
- :param scale: Scaling factor to resize the image.
:param confidence_score: Whether to compute confidence score.
:param attention_map: Whether to plot the attention map.
:param attention_map_level: Level of objects to extract.
:param attention_map_scale: Scaling factor for the attention map.
:param word_separators: List of word separators.
:param line_separators: List of line separators.
- :param image_max_width: Resize image
:param predict_objects: Whether to extract objects.
:param threshold_method: Thresholding method. Should be in ["otsu", "simple"].
:param threshold_value: Thresholding value to use for the "simple" thresholding method.
@@ -422,13 +412,14 @@ def run(
device = f"cuda{cuda_device}" if torch.cuda.is_available() else "cpu"
dan_model = DAN(device, temperature)
dan_model.load(model, parameters, charset, mode="eval")
- if image:
+
+ images = image_dir.rglob(f"*{image_extension}") if not image else [image]
+ for image_name in images:
process_image(
- image,
+ image_name,
dan_model,
device,
output,
- scale,
confidence_score,
confidence_score_levels,
attention_map,
@@ -436,28 +427,7 @@ def run(
attention_map_scale,
word_separators,
line_separators,
- image_max_width,
predict_objects,
threshold_method,
threshold_value,
)
- else:
- for image_name in image_dir.rglob(f"*{image_extension}"):
- process_image(
- image_name,
- dan_model,
- device,
- output,
- scale,
- confidence_score,
- confidence_score_levels,
- attention_map,
- attention_map_level,
- attention_map_scale,
- word_separators,
- line_separators,
- image_max_width,
- predict_objects,
- threshold_method,
- threshold_value,
- )
diff --git a/dan/transforms.py b/dan/transforms.py
index 51bc98280eb3c674db6f12847908efb4a4d71d20..f17aa90038bcb1c75f2d0887cf91e487d292225a 100644
--- a/dan/transforms.py
+++ b/dan/transforms.py
@@ -178,7 +178,9 @@ class DPIAdjusting(ImageOnlyTransform):
return np.array(augmented_image)
-def get_preprocessing_transforms(preprocessings: list) -> Compose:
+def get_preprocessing_transforms(
+ preprocessings: list, to_pil_image: bool = False
+) -> Compose:
"""
Returns a list of transformations to be applied to the image.
"""
@@ -198,7 +200,8 @@ def get_preprocessing_transforms(preprocessings: list) -> Compose:
)
case Preprocessing.FixedWidthResize:
transforms.append(FixedWidthResize(width=preprocessing["fixed_width"]))
- transforms.append(ToPILImage())
+ if to_pil_image:
+ transforms.append(ToPILImage())
return Compose(transforms)
diff --git a/dan/utils.py b/dan/utils.py
index 822ae13d2c3f53048655cbb29332c895f70fb2c1..98820b0937fbbde158d063389cd5c1cbcb497ba2 100644
--- a/dan/utils.py
+++ b/dan/utils.py
@@ -1,6 +1,6 @@
# -*- coding: utf-8 -*-
-import cv2
import torch
+import torchvision.io as torchvision
# Layout begin-token to end-token
SEM_MATCHING_TOKENS = {"ⓘ": "â’¾", "â““": "â’¹", "â“¢": "Ⓢ", "â“’": "â’¸", "â“Ÿ": "â“…", "â“": "â’¶"}
@@ -43,18 +43,13 @@ def pad_images(images):
return padded_images
-def read_image(filename, scale=1.0):
+def read_image(path):
"""
- Read image and rescale it
- :param filename: Image path
- :param scale: Scaling factor before prediction
+ Read image with torch
+ :param path: Path of the image to load.
"""
- image = cv2.cvtColor(cv2.imread(str(filename)), cv2.COLOR_BGR2RGB)
- if scale != 1.0:
- width = int(image.shape[1] * scale)
- height = int(image.shape[0] * scale)
- image = cv2.resize(image, (width, height), interpolation=cv2.INTER_AREA)
- return image
+ img = torchvision.read_image(path, mode=torchvision.ImageReadMode.RGB)
+ return img.to(dtype=torch.get_default_dtype()).div(255)
# Charset / labels conversion
diff --git a/docs/get_started/training.md b/docs/get_started/training.md
index 8eb5dc0c2c08b8f63d489cedfce0fc8962b717f6..71f672d79a3a82efdab223ecf04e20c8f2292c58 100644
--- a/docs/get_started/training.md
+++ b/docs/get_started/training.md
@@ -70,5 +70,11 @@ parameters:
dec_num_heads: int
dec_att_dropout: float
dec_res_dropout: float
+ preprocessings:
+ - type: str
+ max_height: int
+ max_width: int
+ fixed_height: int
+ fixed_width: int
```
2. Apply a trained DAN model on an image using the [predict command](../usage/predict.md).
diff --git a/docs/usage/predict.md b/docs/usage/predict.md
index 0532c603a85179d98345c59a279cb85a5d3db425..98da7a254f18c7485bba234a4dbabf02db0d45c3 100644
--- a/docs/usage/predict.md
+++ b/docs/usage/predict.md
@@ -13,7 +13,6 @@ Use the `teklia-dan predict` command to apply a trained DAN model on an image.
| `--parameters` | Path to the YAML parameters file. | `Path` | |
| `--charset` | Path to the charset file. | `Path` | |
| `--output` | Path to the output folder. Results will be saved in this directory. | `Path` | |
-| `--scale` | Image scaling factor before feeding it to DAN. | `float` | `1.0` |
| `--confidence-score` | Whether to return confidence scores. | `bool` | `False` |
| `--confidence-score-levels` | Level to return confidence scores. Should be any combination of `["line", "word", "char"]`. | `str` | |
| `--attention-map` | Whether to plot attention maps. | `bool` | `False` |
@@ -37,7 +36,6 @@ teklia-dan predict \
--parameters dan_humu_page/parameters.yml \
--charset dan_humu_page/charset.pkl \
--output dan_humu_page/predict/ \
- --scale 0.5 \
--confidence-score
```
It will create the following JSON file named `dan_humu_page/predict/example.json`
@@ -60,7 +58,6 @@ teklia-dan predict \
--parameters dan_humu_page/parameters.yml \
--charset dan_humu_page/charset.pkl \
--output dan_humu_page/predict/ \
- --scale 0.5 \
--confidence-score \
--attention-map \
```
@@ -88,7 +85,6 @@ teklia-dan predict \
--parameters dan_humu_page/parameters.yml \
--charset dan_humu_page/charset.pkl \
--output dan_humu_page/predict/ \
- --scale 0.5 \
--confidence-score \
--attention-map \
--attention-map-level word \
@@ -118,7 +114,6 @@ teklia-dan predict \
--parameters dan_humu_page/parameters.yml \
--charset dan_humu_page/charset.pkl \
--output dan_humu_page/predict/ \
- --scale 0.5 \
--attention-map \
--predict-objects \
--threshold-method otsu
diff --git a/tests/data/prediction/parameters.yml b/tests/data/prediction/parameters.yml
index 76f665e2b0354dfbdf3122e0482439703fdc82da..dbb4bc0f0779e59f14248286ad42586fc1520345 100644
--- a/tests/data/prediction/parameters.yml
+++ b/tests/data/prediction/parameters.yml
@@ -22,3 +22,7 @@ parameters:
dec_num_heads: 4
dec_att_dropout: 0.1
dec_res_dropout: 0.1
+ preprocessings:
+ - type: "max_resize"
+ max_height: 1500
+ max_width: 1500
diff --git a/tests/test_prediction.py b/tests/test_prediction.py
index 1a920275a0ba53024e2c61a67032c31499bdee66..ae289d234d71b802c88be6f3b8baebdf496c43cc 100644
--- a/tests/test_prediction.py
+++ b/tests/test_prediction.py
@@ -3,11 +3,9 @@
import json
import pytest
-import torch
from dan.predict.prediction import DAN
from dan.predict.prediction import run as run_prediction
-from dan.utils import read_image
@pytest.mark.parametrize(
@@ -46,12 +44,12 @@ def test_predict(
mode="eval",
)
- image = read_image(prediction_data_path / "images" / image_name)
- image = dan_model.preprocess(image)
+ image_path = prediction_data_path / "images" / image_name
+ image = dan_model.preprocess(str(image_path))
- input_tensor = torch.tensor(image).permute(2, 0, 1).unsqueeze(0)
+ input_tensor = image.unsqueeze(0)
input_tensor = input_tensor.to(device)
- input_sizes = [image.shape[:2]]
+ input_sizes = [image.shape[1:]]
prediction = dan_model.predict(input_tensor, input_sizes)
@@ -259,7 +257,6 @@ def test_run_prediction(
parameters=prediction_data_path / "parameters.yml",
charset=prediction_data_path / "charset.pkl",
output=tmp_path,
- scale=1,
confidence_score=True if confidence_score else False,
confidence_score_levels=confidence_score if confidence_score else [],
attention_map=False,
@@ -268,7 +265,6 @@ def test_run_prediction(
word_separators=[" ", "\n"],
line_separators=["\n"],
temperature=temperature,
- image_max_width=None,
predict_objects=False,
threshold_method="otsu",
threshold_value=0,