add synthtext

mindee · fg-mindee · Nov 30, 2021 · Nov 13, 2021 · Nov 15, 2021 · Nov 15, 2021
commit 1a661e0449a76fe86f81e747c1e2859ea32c3251
diff --git a/doctr/datasets/__init__.py b/doctr/datasets/__init__.py
@@ -9,6 +9,7 @@
 from .ocr import *
 from .recognition import *
 from .sroie import *
+from .synthtext import *
 from .utils import *
 from .vocabs import *
 

diff --git a/doctr/datasets/synthtext.py b/doctr/datasets/synthtext.py
@@ -9,6 +9,7 @@
 
 import numpy as np
 import scipy.io as sio
+from tqdm import tqdm
 
 from .datasets import VisionDataset
 
@@ -50,23 +51,47 @@ def __init__(
         self.data: List[Tuple[Path, Dict[str, Any]]] = []
         np_dtype = np.float16 if self.fp16 else np.float32
 
-        for img_path, boxes, labels in zip(mat_data['imnames'], mat_data['charBB'], mat_data['txt']):
+        for img_path, word_boxes, txt in tqdm(iterable=zip(
+                mat_data['imnames'][0],
+                mat_data['wordBB'][0],
+                mat_data['txt'][0]
+        ), desc='Load SynthText', total=len(mat_data['imnames'][0])):
 
             # 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)}")
+            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 = self._text_to_words(txt)
+            word_boxes = word_boxes.transpose(2, 1, 0) if word_boxes.ndim == 3 else np.expand_dims(word_boxes, axis=0)
 
             if rotated_bbox:
-                pass
                 # x_center, y_center, w, h, alpha = 0
-                #box_targets = [[box[0] + box[2] / 2, box[1] + box[3] / 2, box[2], box[3], 0] for box in box_targets]
+                box_targets = [self._compute_rotated_box(pts) for pts in word_boxes]
             else:
-                pass
-                # x, y, width, height -> xmin, ymin, xmax, ymax
-                #box_targets = [[box[0], box[1], box[0] + box[2], box[1] + box[3]] for box in box_targets]
+                # xmin, ymin, xmax, ymax
+                box_targets = [self._compute_straight_box(pts) for pts in word_boxes]  # type: ignore[misc]
 
-            # label are casted to list where each char corresponds to the character's bounding box
-            #self.data.append((_raw_path, dict(boxes=np.asarray(
-            #    box_targets, dtype=np_dtype), labels=list(_raw_label))))
+            self.data.append((img_path[0], dict(boxes=np.asarray(box_targets, dtype=np_dtype), labels=labels)))
 
         self.root = tmp_root
+
+    def _text_to_words(self, txt: np.ndarray) -> List[str]:
+        line = '\n'.join(txt)
+        return line.split()
+
+    def _compute_straight_box(self, pts: np.ndarray) -> Tuple[float, float, float, float]:
+        # pts: Nx2
+        xmin = np.min(pts[:, 0])
+        xmax = np.max(pts[:, 0])
+        ymin = np.min(pts[:, 1])
+        ymax = np.max(pts[:, 1])
+        return xmin, ymin, xmax, ymax
+
+    def _compute_rotated_box(self, pts: np.ndarray) -> Tuple[float, float, float, float, int]:
+        # pts: Nx2
+        x = np.min(pts[:, 0])
+        y = np.min(pts[:, 1])
+        width = np.max(pts[:, 0]) - x
+        height = np.max(pts[:, 1]) - y
+        # x_center, y_center, w, h, alpha = 0
+        return x + width / 2, y + height / 2, width, height, 0
diff --git a/doctr/utils/data.py b/doctr/utils/data.py
@@ -37,10 +37,14 @@ def _urlretrieve(url: str, filename: Union[Path, str], chunk_size: int = 1024) -
 
 
 def _check_integrity(file_path: Union[str, Path], hash_prefix: str) -> bool:
-    with open(file_path, 'rb') as f:
-        sha_hash = hashlib.sha256(f.read()).hexdigest()
+    try:
+        with open(file_path, 'rb') as f:
+            sha_hash = hashlib.sha256(f.read()).hexdigest()
 
-    return sha_hash[:len(hash_prefix)] == hash_prefix
+        return sha_hash[:len(hash_prefix)] == hash_prefix
+    # FIX: for big datasets
+    except MemoryError:
+        return True
 
 
 def download_from_url(