Spaces:
Runtime error
Runtime error
File size: 2,550 Bytes
153628e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
# Copyright (C) 2021-2024, 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, Dict, List, Tuple
import numpy as np
from .datasets import AbstractDataset
__all__ = ["OCRDataset"]
class OCRDataset(AbstractDataset):
"""Implements an OCR dataset
>>> from doctr.datasets import OCRDataset
>>> train_set = OCRDataset(img_folder="/path/to/images",
>>> label_file="/path/to/labels.json")
>>> img, target = train_set[0]
Args:
----
img_folder: local path to image folder (all jpg at the root)
label_file: local path to the label file
use_polygons: whether polygons should be considered as rotated bounding box (instead of straight ones)
**kwargs: keyword arguments from `AbstractDataset`.
"""
def __init__(
self,
img_folder: str,
label_file: str,
use_polygons: bool = False,
**kwargs: Any,
) -> None:
super().__init__(img_folder, **kwargs)
# List images
self.data: List[Tuple[str, Dict[str, Any]]] = []
np_dtype = np.float32
with open(label_file, "rb") as f:
data = json.load(f)
for img_name, annotations in data.items():
# Get image path
img_name = Path(img_name)
# File existence check
if not os.path.exists(os.path.join(self.root, img_name)):
raise FileNotFoundError(f"unable to locate {os.path.join(self.root, img_name)}")
# handle empty images
if len(annotations["typed_words"]) == 0:
self.data.append((img_name, dict(boxes=np.zeros((0, 4), dtype=np_dtype), labels=[])))
continue
# Unpack the straight boxes (xmin, ymin, xmax, ymax)
geoms = [list(map(float, obj["geometry"][:4])) for obj in annotations["typed_words"]]
if use_polygons:
# (x, y) coordinates of top left, top right, bottom right, bottom left corners
geoms = [
[geom[:2], [geom[2], geom[1]], geom[2:], [geom[0], geom[3]]] # type: ignore[list-item]
for geom in geoms
]
text_targets = [obj["value"] for obj in annotations["typed_words"]]
self.data.append((img_name, dict(boxes=np.asarray(geoms, dtype=np_dtype), labels=text_targets)))
|