update

mindee · charlesmindee · Apr 28, 2022 · Jan 11, 2022 · Jan 11, 2022 · Jan 18, 2022
commit 7877ac5897df1cd5d3707e496fbbb9f8f21c51ab
diff --git a/tests/pytorch/test_datasets_pt.py b/tests/pytorch/test_datasets_pt.py
@@ -10,22 +10,21 @@
 from doctr.transforms import Resize
 
 
-def _validate_dataset(ds, input_size, batch_size=2, class_indices=False, is_polygons=False, target_includes_boxes=True):
+def _validate_dataset(ds, input_size, batch_size=2, class_indices=False, is_polygons=False):
 
     # Fetch one sample
     img, target = ds[0]
     assert isinstance(img, torch.Tensor)
     assert img.shape == (3, *input_size)
     assert img.dtype == torch.float32
     assert isinstance(target, dict)
-    if target_includes_boxes:
-        assert isinstance(target['boxes'], np.ndarray) and target['boxes'].dtype == np.float32
-        if is_polygons:
-            assert target['boxes'].ndim == 3 and target['boxes'].shape[1:] == (4, 2)
-        else:
-            assert target['boxes'].ndim == 2 and target['boxes'].shape[1:] == (4,)
-        assert np.all(np.logical_and(target['boxes'] <= 1, target['boxes'] >= 0))
-        assert len(target['labels']) == len(target['boxes'])
+    assert isinstance(target['boxes'], np.ndarray) and target['boxes'].dtype == np.float32
+    if is_polygons:
+        assert target['boxes'].ndim == 3 and target['boxes'].shape[1:] == (4, 2)
+    else:
+        assert target['boxes'].ndim == 2 and target['boxes'].shape[1:] == (4,)
+    assert np.all(np.logical_and(target['boxes'] <= 1, target['boxes'] >= 0))
+    assert len(target['labels']) == len(target['boxes'])
     if class_indices:
         assert isinstance(target['labels'], np.ndarray) and target['labels'].dtype == np.int64
     else:

diff --git a/tests/tensorflow/test_datasets_tf.py b/tests/tensorflow/test_datasets_tf.py
@@ -10,22 +10,21 @@
 from doctr.transforms import Resize
 
 
-def _validate_dataset(ds, input_size, batch_size=2, class_indices=False, is_polygons=False, target_includes_boxes=True):
+def _validate_dataset(ds, input_size, batch_size=2, class_indices=False, is_polygons=False):
 
     # Fetch one sample
     img, target = ds[0]
     assert isinstance(img, tf.Tensor)
     assert img.shape == (*input_size, 3)
     assert img.dtype == tf.float32
     assert isinstance(target, dict)
-    if target_includes_boxes:
-        assert isinstance(target['boxes'], np.ndarray) and target['boxes'].dtype == np.float32
-        if is_polygons:
-            assert target['boxes'].ndim == 3 and target['boxes'].shape[1:] == (4, 2)
-        else:
-            assert target['boxes'].ndim == 2 and target['boxes'].shape[1:] == (4,)
-        assert np.all(np.logical_and(target['boxes'] <= 1, target['boxes'] >= 0))
-        assert len(target['labels']) == len(target['boxes'])
+    assert isinstance(target['boxes'], np.ndarray) and target['boxes'].dtype == np.float32
+    if is_polygons:
+        assert target['boxes'].ndim == 3 and target['boxes'].shape[1:] == (4, 2)
+    else:
+        assert target['boxes'].ndim == 2 and target['boxes'].shape[1:] == (4,)
+    assert np.all(np.logical_and(target['boxes'] <= 1, target['boxes'] >= 0))
+    assert len(target['labels']) == len(target['boxes'])
     if class_indices:
         assert isinstance(target['labels'], np.ndarray) and target['labels'].dtype == np.int64
     else: