add repr

mindee · Nov 15, 2021 · 09d0d83 · 09d0d83
1 parent 9099e95
commit 09d0d83
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/doctr/datasets/synthtext.py b/doctr/datasets/synthtext.py
@@ -45,6 +45,7 @@ def __init__(
 
         super().__init__(url=self.URL, file_name=None, file_hash=self.SHA256, extract_archive=True, **kwargs)
         self.sample_transforms = sample_transforms
+        self.train = train
 
         # Load mat data
         tmp_root = os.path.join(self.root, 'SynthText')
@@ -77,9 +78,11 @@ def __init__(
 
             self.data.append((img_path[0], dict(boxes=np.asarray(box_targets, dtype=np_dtype), labels=labels)))
 
-        if train:
+        if self.train:
             self.data = self.data[:int(len(self.data) * 0.9)]
         else:
             self.data = self.data[int(len(self.data) * 0.9):]
-
         self.root = tmp_root
+
+    def extra_repr(self) -> str:
+        return f"train={self.train}"