refactor: Refactored dataset constructors

mindee · Jun 5, 2021 · ea704e9 · ea704e9
1 parent f387d79
commit ea704e9
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/doctr/datasets/detection.py b/doctr/datasets/detection.py
@@ -47,12 +47,12 @@ def __init__(
             with open(os.path.join(label_folder, img_path + '.json'), 'rb') as f:
                 boxes = json.load(f)
             bboxes = np.asarray(boxes["boxes_1"] + boxes["boxes_2"] + boxes["boxes_3"], dtype=np.float32)
-            if not rotated_bbox:
-                # Switch to xmin, ymin, xmax, ymax
-                bboxes = np.concatenate((bboxes.min(axis=1), bboxes.max(axis=1)), axis=1)
-            else:
+            if rotated_bbox:
                 # Switch to rotated rects
                 bboxes = np.asarray([list(fit_rbbox(box)) for box in bboxes], dtype=np.float32)
+            else:
+                # Switch to xmin, ymin, xmax, ymax
+                bboxes = np.concatenate((bboxes.min(axis=1), bboxes.max(axis=1)), axis=1)
 
             is_ambiguous = [False] * (len(boxes["boxes_1"]) + len(boxes["boxes_2"])) + [True] * len(boxes["boxes_3"])
             self.data.append((img_path, dict(boxes=bboxes, flags=np.asarray(is_ambiguous))))

diff --git a/doctr/datasets/ocr.py b/doctr/datasets/ocr.py
@@ -62,15 +62,14 @@ def __init__(
             for box in file_dic["coordinates"]:
                 if rotated_bbox:
                     x, y, w, h, alpha = fit_rbbox(np.asarray(box, dtype=np.float32))
+                    box = [x, y, w, h, alpha]
                     is_valid.append(w > 0 and h > 0)
-                    if is_valid[-1]:
-                        box_targets.append([x, y, w, h, alpha])
                 else:
                     xs, ys = zip(*box)
                     box = [min(xs), min(ys), max(xs), max(ys)]
                     is_valid.append(box[0] < box[2] and box[1] < box[3])
-                    if is_valid[-1]:
-                        box_targets.append(box)
+                if is_valid[-1]:
+                    box_targets.append(box)
 
             text_targets = [word for word, _valid in zip(file_dic["string"], is_valid) if _valid]
             self.data.append((img_name, dict(boxes=np.asarray(box_targets, dtype=np.float32), labels=text_targets)))