# Copyright (C) 2021-2025, Mindee.
# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
import json
import os
from pathlib import Path
from typing import Any
import numpy as np
from tqdm import tqdm
from .datasets import VisionDataset
from .utils import convert_target_to_relative, crop_bboxes_from_image
__all__ = ["CORD"]
[docs]
class CORD(VisionDataset):
"""CORD dataset from `"CORD: A Consolidated Receipt Dataset forPost-OCR Parsing"
<https://openreview.net/pdf?id=SJl3z659UH>`_.
.. image:: https://doctr-static.mindee.com/models?id=v0.5.0/cord-grid.png&src=0
:align: center
>>> from doctr.datasets import CORD
>>> train_set = CORD(train=True, download=True)
>>> img, target = train_set[0]
Args:
train: whether the subset should be the training one
use_polygons: whether polygons should be considered as rotated bounding box (instead of straight ones)
recognition_task: whether the dataset should be used for recognition task
detection_task: whether the dataset should be used for detection task
**kwargs: keyword arguments from `VisionDataset`.
"""
TRAIN = (
"https://doctr-static.mindee.com/models?id=v0.1.1/cord_train.zip&src=0",
"45f9dc77f126490f3e52d7cb4f70ef3c57e649ea86d19d862a2757c9c455d7f8",
"cord_train.zip",
)
TEST = (
"https://doctr-static.mindee.com/models?id=v0.1.1/cord_test.zip&src=0",
"8c895e3d6f7e1161c5b7245e3723ce15c04d84be89eaa6093949b75a66fb3c58",
"cord_test.zip",
)
def __init__(
self,
train: bool = True,
use_polygons: bool = False,
recognition_task: bool = False,
detection_task: bool = False,
**kwargs: Any,
) -> None:
url, sha256, name = self.TRAIN if train else self.TEST
super().__init__(
url,
name,
sha256,
True,
pre_transforms=convert_target_to_relative if not recognition_task else None,
**kwargs,
)
if recognition_task and detection_task:
raise ValueError(
"`recognition_task` and `detection_task` cannot be set to True simultaneously. "
+ "To get the whole dataset with boxes and labels leave both parameters to False."
)
# list images
tmp_root = os.path.join(self.root, "image")
self.data: list[tuple[str | np.ndarray, str | dict[str, Any] | np.ndarray]] = []
self.train = train
np_dtype = np.float32
for img_path in tqdm(
iterable=os.listdir(tmp_root), desc="Preparing and Loading CORD", total=len(os.listdir(tmp_root))
):
# File existence check
if not os.path.exists(os.path.join(tmp_root, img_path)):
raise FileNotFoundError(f"unable to locate {os.path.join(tmp_root, img_path)}")
stem = Path(img_path).stem
_targets = []
with open(os.path.join(self.root, "json", f"{stem}.json"), "rb") as f:
label = json.load(f)
for line in label["valid_line"]:
for word in line["words"]:
if len(word["text"]) > 0:
x = word["quad"]["x1"], word["quad"]["x2"], word["quad"]["x3"], word["quad"]["x4"]
y = word["quad"]["y1"], word["quad"]["y2"], word["quad"]["y3"], word["quad"]["y4"]
box: list[float] | np.ndarray
if use_polygons:
# (x, y) coordinates of top left, top right, bottom right, bottom left corners
box = np.array(
[
[x[0], y[0]],
[x[1], y[1]],
[x[2], y[2]],
[x[3], y[3]],
],
dtype=np_dtype,
)
else:
# Reduce 8 coords to 4 -> xmin, ymin, xmax, ymax
box = [min(x), min(y), max(x), max(y)]
_targets.append((word["text"], box))
text_targets, box_targets = zip(*_targets)
if recognition_task:
crops = crop_bboxes_from_image(
img_path=os.path.join(tmp_root, img_path), geoms=np.asarray(box_targets, dtype=int).clip(min=0)
)
for crop, label in zip(crops, list(text_targets)):
if " " not in label:
self.data.append((crop, label))
elif detection_task:
self.data.append((img_path, np.asarray(box_targets, dtype=int).clip(min=0)))
else:
self.data.append((
img_path,
dict(boxes=np.asarray(box_targets, dtype=int).clip(min=0), labels=list(text_targets)),
))
self.root = tmp_root
def extra_repr(self) -> str:
return f"train={self.train}"