Train custom KIE

[lucas2lab]

Hello,

I am trying to train a custom KIE as I would like to recognise some specific entities on the documents I am scanning and processing with doctr.

To do so I proceeded as follows:

1. used a vanilla model to process a sample of documents
2. saved the bboxes and uploaded them to roboflow where I cleaned them up (removed wrong boxes, refined boxes location) and annotated as new classes the entities that interest me
3. then, I exported the dataset and transformed it into a doctr one as per this documentation https://github.com/mindee/doctr/tree/main/references/detection I kept a word class for all basic words and created two new classes for my custom ones. I believe those will be returned by the KIE?
4. then I trained a model using references/detection/train_pytorch.py
5. I am now trying to use this model to do some KIE but encounter some error.

Here is an extract of the dataset in JSON

{
  "5_jpg.rf.fd99a4978c094183d448bfbcfb6e8a3b.jpg": {
    "img_dimensions": [
      3368,
      2380
    ],
    "img_hash": "39f1664ad668d6aa39cfecd64996fa1086094012dcd76805514eba077231e661",
    "polygons": {
      "custom_class1": [
        [
          [
            1936,
            121
          ],
          [
            1999,
            121
          ],
          [
            1999,
            161
          ],
          [
            1936,
            161
          ]
        ],
        [
          [
            183,
            167
          ],
          [
            249,
            167
          ],
          [
            249,
            210
          ],
          [
            183,
            210
          ]
        ],

Install deps

!pip -q install tf2onnx
!git clone --depth=1 https://github.com/mindee/doctr.git
!pip -q install -e doctr/.[torch]
!pip install -q matplotlib mplcursors
!cd doctr && pip -q install -r references/requirements.txt
!apt install -y fonts-freefont-ttf

Run training script

!python doctr/references/detection/train_pytorch.py /content/AWB-OCR-Detector-1-doctr/train /content/AWB-OCR-Detector-1-doctr/valid db_resnet50 --epochs 5 --device 0 --pretrained --freeze-backbone

Namespace(train_path='/content/AWB-OCR-Detector-1-doctr/train', val_path='/content/AWB-OCR-Detector-1-doctr/valid', arch='db_resnet50', name=None, epochs=5, batch_size=2, device=0, save_interval_epoch=False, input_size=1024, lr=0.001, weight_decay=0, workers=None, resume=None, test_only=False, freeze_backbone=True, show_samples=False, wb=False, push_to_hub=False, pretrained=True, rotation=False, eval_straight=False, sched='poly', amp=False, find_lr=False, early_stop=False, early_stop_epochs=5, early_stop_delta=0.01)
Validation set loaded in 0.248s (6 samples in 3 batches)
Downloading https://doctr-static.mindee.com/models?id=v0.7.0/db_resnet50-79bd7d70.pt&src=0 to /root/.cache/doctr/models/db_resnet50-79bd7d70.pt
102022144it [00:01, 52250613.63it/s]                  
Train set loaded in 0.03282s (17 samples in 8 batches)

  0% 0/8 [00:00<?, ?it/s]
Training loss: 2.91673:   0% 0/8 [00:27<?, ?it/s]
Training loss: 2.91673:  12% 1/8 [00:27<03:10, 27.15s/it]
Training loss: 2.83323:  12% 1/8 [00:31<03:10, 27.15s/it]
Training loss: 2.83323:  25% 2/8 [00:31<01:22, 13.71s/it]
Training loss: 2.77549:  25% 2/8 [00:34<01:22, 13.71s/it]
Training loss: 2.77549:  38% 3/8 [00:34<00:45,  9.07s/it]
Training loss: 2.70702:  38% 3/8 [00:37<00:45,  9.07s/it]
Training loss: 2.70702:  50% 4/8 [00:37<00:26,  6.55s/it]
Training loss: 2.66043:  50% 4/8 [00:40<00:26,  6.55s/it]
Training loss: 2.66043:  62% 5/8 [00:40<00:15,  5.19s/it]
Training loss: 2.60811:  62% 5/8 [00:45<00:15,  5.19s/it]
Training loss: 2.60811:  75% 6/8 [00:45<00:10,  5.18s/it]
Training loss: 2.55209:  75% 6/8 [00:47<00:10,  5.18s/it]
Training loss: 2.55209:  88% 7/8 [00:47<00:04,  4.25s/it]
Training loss: 2.51749:  88% 7/8 [00:49<00:04,  4.25s/it]
Training loss: 2.51749: 100% 8/8 [00:49<00:00,  6.18s/it]
100% 3/3 [00:14<00:00,  4.93s/it]
Validation loss decreased inf --> 2.46541: saving state...
Epoch 1/5 - Validation loss: 2.46541 (Recall: 0.00% | Precision: 0.00% | Mean IoU: 0.00%)

NB: despite number of epoch I get Recall: 0.00% | Precision: 0.00% | Mean IoU: 0.00% :(

Load and run an inference

import os
os.environ['USE_TORCH'] = '1'

import torch
from doctr.models import kie_predictor, db_resnet50

det_model = db_resnet50(pretrained=False, pretrained_backbone=False, class_names=['word', 'class1', 'class2'])
det_params = torch.load('/content/db_resnet50_20240729-084124.pt', map_location="cpu")
det_model.load_state_dict(det_params)
model = kie_predictor(det_arch=det_model, reco_arch='crnn_vgg16_bn', pretrained=True)

from doctr.io import DocumentFile

doc = DocumentFile.from_images("/path/to/my/img.jpg")
# Analyze
result = model(doc)

I get the following stacktrace

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
[<ipython-input-6-6c5c9789fd9d>](https://localhost:8080/#) in <cell line: 5>()
      3 doc = DocumentFile.from_images("/xxx.jpg")
      4 # Analyze
----> 5 result = model(doc)

6 frames
[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1530             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531         else:
-> 1532             return self._call_impl(*args, **kwargs)
   1533 
   1534     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1539                 or _global_backward_pre_hooks or _global_backward_hooks
   1540                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541             return forward_call(*args, **kwargs)
   1542 
   1543         try:

[/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py](https://localhost:8080/#) in decorate_context(*args, **kwargs)
    113     def decorate_context(*args, **kwargs):
    114         with ctx_factory():
--> 115             return func(*args, **kwargs)
    116 
    117     return decorate_context

[/content/doctr/doctr/models/kie_predictor/pytorch.py](https://localhost:8080/#) in forward(self, pages, **kwargs)
    103             loc_preds = self.det_predictor(pages, **kwargs)
    104 
--> 105         dict_loc_preds: Dict[str, List[np.ndarray]] = invert_data_structure(loc_preds)  # type: ignore[assignment]
    106 
    107         # Detach objectness scores from loc_preds

[/content/doctr/doctr/models/_utils.py](https://localhost:8080/#) in invert_data_structure(x)
    196         return [dict(zip(x, t)) for t in zip(*x.values())]
    197     elif isinstance(x, list):
--> 198         return {k: [dic[k] for dic in x] for k in x[0]}
    199     else:
    200         raise TypeError(f"Expected input to be either a dict or a list, got {type(input)} instead.")

[/content/doctr/doctr/models/_utils.py](https://localhost:8080/#) in <dictcomp>(.0)
    196         return [dict(zip(x, t)) for t in zip(*x.values())]
    197     elif isinstance(x, list):
--> 198         return {k: [dic[k] for dic in x] for k in x[0]}
    199     else:
    200         raise TypeError(f"Expected input to be either a dict or a list, got {type(input)} instead.")

[/content/doctr/doctr/models/_utils.py](https://localhost:8080/#) in <listcomp>(.0)
    196         return [dict(zip(x, t)) for t in zip(*x.values())]
    197     elif isinstance(x, list):
--> 198         return {k: [dic[k] for dic in x] for k in x[0]}
    199     else:
    200         raise TypeError(f"Expected input to be either a dict or a list, got {type(input)} instead.")

KeyError: 'class1'

Is it a bug on doctr or am I doing something wrong ?

NB: sample size is small as I want to validate the approach before I start to annotate a large quantity of documents.

[jbesq-ml]

Hi @lucas2lab,

Did you have this issue when trying to train? How did you fix it? I would also like to give your thread a bump. I used similar data.

Namespace(train_path='/home/PycharmProjects/doctr', val_path='/home/PycharmProjects/doctr', arch='db_resnet50', name=None, epochs=5, batch_size=2, device=0, save_interval_epoch=False, input_size=1024, lr=0.001, weight_decay=0, workers=None, resume=None, test_only=False, freeze_backbone=False, show_samples=False, wb=False, push_to_hub=False, pretrained=False, rotation=False, eval_straight=False, sched='poly', amp=False, find_lr=False, early_stop=False, early_stop_epochs=5, early_stop_delta=0.01)

File "/home/anaconda3/envs/doctr/lib/python3.9/site-packages/doctr/datasets/detection.py", line 93, in format_polygons

geoms = _polygons if use_polygons else np.concatenate((_polygons.min(axis=1), _polygons.max(axis=1)), axis=1)
numpy.exceptions.AxisError: axis 1 is out of bounds for array of dimension 1