clip box values and add shape and label check

mindee · Jun 17, 2022 · abac30c · abac30c
1 parent 210ecc4
commit abac30c
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/doctr/datasets/svt.py b/doctr/datasets/svt.py
@@ -93,7 +93,8 @@ def __init__(
             if recognition_task:
                 crops = crop_bboxes_from_image(img_path=os.path.join(tmp_root, name.text), geoms=boxes)
                 for crop, label in zip(crops, labels):
-                    self.data.append((crop, dict(labels=[label])))
+                    if crop.shape[0] > 0 and crop.shape[1] > 0 and len(label) > 0:
+                        self.data.append((crop, dict(labels=[label])))
             else:
                 # Convert coordinates to relative
                 w, h = int(resolution.attrib['x']), int(resolution.attrib['y'])
@@ -104,7 +105,7 @@ def __init__(
                     boxes[:, [0, 2]] /= w
                     boxes[:, [1, 3]] /= h
 
-                self.data.append((name.text, dict(boxes=boxes, labels=labels)))
+                self.data.append((name.text, dict(boxes=boxes.clip(0, 1), labels=labels)))
 
         self.root = tmp_root