Spaces:
Runtime error
Runtime error
File size: 6,705 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
# 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 glob
import json
import os
from pathlib import Path
from typing import Any, Dict, List, Tuple, Union
import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm
from .datasets import AbstractDataset
from .utils import convert_target_to_relative, crop_bboxes_from_image
__all__ = ["IMGUR5K"]
class IMGUR5K(AbstractDataset):
"""IMGUR5K dataset from `"TextStyleBrush: Transfer of Text Aesthetics from a Single Example"
<https://arxiv.org/abs/2106.08385>`_ |
`repository <https://github.com/facebookresearch/IMGUR5K-Handwriting-Dataset>`_.
.. image:: https://doctr-static.mindee.com/models?id=v0.5.0/imgur5k-grid.png&src=0
:align: center
:width: 630
:height: 400
>>> # NOTE: You need to download/generate the dataset from the repository.
>>> from doctr.datasets import IMGUR5K
>>> train_set = IMGUR5K(train=True, img_folder="/path/to/IMGUR5K-Handwriting-Dataset/images",
>>> label_path="/path/to/IMGUR5K-Handwriting-Dataset/dataset_info/imgur5k_annotations.json")
>>> img, target = train_set[0]
>>> test_set = IMGUR5K(train=False, img_folder="/path/to/IMGUR5K-Handwriting-Dataset/images",
>>> label_path="/path/to/IMGUR5K-Handwriting-Dataset/dataset_info/imgur5k_annotations.json")
>>> img, target = test_set[0]
Args:
----
img_folder: folder with all the images of the dataset
label_path: path to the annotations file of the dataset
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
**kwargs: keyword arguments from `AbstractDataset`.
"""
def __init__(
self,
img_folder: str,
label_path: str,
train: bool = True,
use_polygons: bool = False,
recognition_task: bool = False,
**kwargs: Any,
) -> None:
super().__init__(
img_folder, pre_transforms=convert_target_to_relative if not recognition_task else None, **kwargs
)
# File existence check
if not os.path.exists(label_path) or not os.path.exists(img_folder):
raise FileNotFoundError(f"unable to locate {label_path if not os.path.exists(label_path) else img_folder}")
self.data: List[Tuple[Union[str, Path, np.ndarray], Union[str, Dict[str, Any]]]] = []
self.train = train
np_dtype = np.float32
img_names = os.listdir(img_folder)
train_samples = int(len(img_names) * 0.9)
set_slice = slice(train_samples) if self.train else slice(train_samples, None)
# define folder to write IMGUR5K recognition dataset
reco_folder_name = "IMGUR5K_recognition_train" if self.train else "IMGUR5K_recognition_test"
reco_folder_name = "Poly_" + reco_folder_name if use_polygons else reco_folder_name
reco_folder_path = os.path.join(os.path.dirname(self.root), reco_folder_name)
reco_images_counter = 0
if recognition_task and os.path.isdir(reco_folder_path):
self._read_from_folder(reco_folder_path)
return
elif recognition_task and not os.path.isdir(reco_folder_path):
os.makedirs(reco_folder_path, exist_ok=False)
with open(label_path) as f:
annotation_file = json.load(f)
for img_name in tqdm(iterable=img_names[set_slice], desc="Unpacking IMGUR5K", total=len(img_names[set_slice])):
img_path = Path(img_folder, img_name)
img_id = img_name.split(".")[0]
# 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)}")
# some files have no annotations which are marked with only a dot in the 'word' key
# ref: https://github.com/facebookresearch/IMGUR5K-Handwriting-Dataset/blob/main/README.md
if img_id not in annotation_file["index_to_ann_map"].keys():
continue
ann_ids = annotation_file["index_to_ann_map"][img_id]
annotations = [annotation_file["ann_id"][a_id] for a_id in ann_ids]
labels = [ann["word"] for ann in annotations if ann["word"] != "."]
# x_center, y_center, width, height, angle
_boxes = [
list(map(float, ann["bounding_box"].strip("[ ]").split(", ")))
for ann in annotations
if ann["word"] != "."
]
# (x, y) coordinates of top left, top right, bottom right, bottom left corners
box_targets = [cv2.boxPoints(((box[0], box[1]), (box[2], box[3]), box[4])) for box in _boxes] # type: ignore[arg-type]
if not use_polygons:
# xmin, ymin, xmax, ymax
box_targets = [np.concatenate((points.min(0), points.max(0)), axis=-1) for points in box_targets]
# filter images without boxes
if len(box_targets) > 0:
if recognition_task:
crops = crop_bboxes_from_image(
img_path=os.path.join(self.root, img_name), geoms=np.asarray(box_targets, dtype=np_dtype)
)
for crop, label in zip(crops, labels):
if crop.shape[0] > 0 and crop.shape[1] > 0 and len(label) > 0:
# write data to disk
with open(os.path.join(reco_folder_path, f"{reco_images_counter}.txt"), "w") as f:
f.write(label)
tmp_img = Image.fromarray(crop)
tmp_img.save(os.path.join(reco_folder_path, f"{reco_images_counter}.png"))
reco_images_counter += 1
else:
self.data.append((img_path, dict(boxes=np.asarray(box_targets, dtype=np_dtype), labels=labels)))
if recognition_task:
self._read_from_folder(reco_folder_path)
def extra_repr(self) -> str:
return f"train={self.train}"
def _read_from_folder(self, path: str) -> None:
for img_path in glob.glob(os.path.join(path, "*.png")):
with open(os.path.join(path, f"{os.path.basename(img_path)[:-4]}.txt"), "r") as f:
self.data.append((img_path, f.read()))
|