Spaces:
Runtime error
Runtime error
File size: 5,387 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 |
# 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 os
from typing import Any, Dict, List, Tuple, Union
import numpy as np
from PIL import Image
from scipy import io as sio
from tqdm import tqdm
from .datasets import VisionDataset
from .utils import convert_target_to_relative, crop_bboxes_from_image
__all__ = ["SynthText"]
class SynthText(VisionDataset):
"""SynthText dataset from `"Synthetic Data for Text Localisation in Natural Images"
<https://arxiv.org/abs/1604.06646>`_ | `"repository" <https://github.com/ankush-me/SynthText>`_ |
`"website" <https://www.robots.ox.ac.uk/~vgg/data/scenetext/>`_.
.. image:: https://doctr-static.mindee.com/models?id=v0.5.0/svt-grid.png&src=0
:align: center
>>> from doctr.datasets import SynthText
>>> train_set = SynthText(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
**kwargs: keyword arguments from `VisionDataset`.
"""
URL = "https://thor.robots.ox.ac.uk/~vgg/data/scenetext/SynthText.zip"
SHA256 = "28ab030485ec8df3ed612c568dd71fb2793b9afbfa3a9d9c6e792aef33265bf1"
def __init__(
self,
train: bool = True,
use_polygons: bool = False,
recognition_task: bool = False,
**kwargs: Any,
) -> None:
super().__init__(
self.URL,
None,
file_hash=None,
extract_archive=True,
pre_transforms=convert_target_to_relative if not recognition_task else None,
**kwargs,
)
self.train = train
self.data: List[Tuple[Union[str, np.ndarray], Union[str, Dict[str, Any]]]] = []
np_dtype = np.float32
# Load mat data
tmp_root = os.path.join(self.root, "SynthText") if self.SHA256 else self.root
# define folder to write SynthText recognition dataset
reco_folder_name = "SynthText_recognition_train" if self.train else "SynthText_recognition_test"
reco_folder_name = "Poly_" + reco_folder_name if use_polygons else reco_folder_name
reco_folder_path = os.path.join(tmp_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)
mat_data = sio.loadmat(os.path.join(tmp_root, "gt.mat"))
train_samples = int(len(mat_data["imnames"][0]) * 0.9)
set_slice = slice(train_samples) if self.train else slice(train_samples, None)
paths = mat_data["imnames"][0][set_slice]
boxes = mat_data["wordBB"][0][set_slice]
labels = mat_data["txt"][0][set_slice]
del mat_data
for img_path, word_boxes, txt in tqdm(
iterable=zip(paths, boxes, labels), desc="Unpacking SynthText", total=len(paths)
):
# File existence check
if not os.path.exists(os.path.join(tmp_root, img_path[0])):
raise FileNotFoundError(f"unable to locate {os.path.join(tmp_root, img_path[0])}")
labels = [elt for word in txt.tolist() for elt in word.split()]
# (x, y) coordinates of top left, top right, bottom right, bottom left corners
word_boxes = (
word_boxes.transpose(2, 1, 0)
if word_boxes.ndim == 3
else np.expand_dims(word_boxes.transpose(1, 0), axis=0)
)
if not use_polygons:
# xmin, ymin, xmax, ymax
word_boxes = np.concatenate((word_boxes.min(axis=1), word_boxes.max(axis=1)), axis=1)
if recognition_task:
crops = crop_bboxes_from_image(img_path=os.path.join(tmp_root, img_path[0]), geoms=word_boxes)
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[0], dict(boxes=np.asarray(word_boxes, dtype=np_dtype), labels=labels)))
if recognition_task:
self._read_from_folder(reco_folder_path)
self.root = tmp_root
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()))
|