Merge pull request #330 from activeloopai/fix/segmentation

Fixed segmentation

hub/schema/segmentation.py

@@ -63,10 +63,14 @@ def get_attr_dict(self):
     def __str__(self):
         out = super().__str__()
         out = "Segmentation" + out[6:-1]
-        out = out + ", names=" + self.names if self.names is not None else out
         out = (
-            out + ", num_classes=" + self.num_classes
-            if self.num_classes is not None
+            out + ", names=" + str(self.class_labels._names)
+            if self.class_labels._names is not None
+            else out
+        )
+        out = (
+            out + ", num_classes=" + str(self.class_labels._num_classes)
+            if self.class_labels._num_classes is not None
             else out
         )
         out += ")"

hub/schema/tests/test_features.py

@@ -1,3 +1,4 @@
+from hub.schema import Segmentation
 from hub.schema.class_label import ClassLabel, _load_names_from_file
 from hub.schema.features import HubSchema, SchemaDict
 import pytest
@@ -49,9 +50,33 @@ def test_feature_dict_repr():
     assert expected_output == feature_dict_object.__repr__()
 
 
+def test_segmentation_repr():
+    seg1 = Segmentation(shape=(3008, 3008), dtype="uint8", num_classes=5)
+    seg2 = Segmentation(
+        shape=(3008, 3008), dtype="uint8", names=["apple", "orange", "banana"]
+    )
+
+    text1 = "Segmentation(shape=(3008, 3008), dtype='uint8', num_classes=5)"
+    text2 = "Segmentation(shape=(3008, 3008), dtype='uint8', names=['apple', 'orange', 'banana'], num_classes=3)"
+    assert seg1.__repr__() == text1
+    assert seg2.__repr__() == text2
+
+
+def test_classlabel_repr():
+    cl1 = ClassLabel(num_classes=5)
+    cl2 = ClassLabel(names=["apple", "orange", "banana"])
+
+    text1 = "ClassLabel(shape=(), dtype='int64', num_classes=5)"
+    text2 = "ClassLabel(shape=(), dtype='int64', names=['apple', 'orange', 'banana'], num_classes=3)"
+    assert cl1.__repr__() == text1
+    assert cl2.__repr__() == text2
+
+
 if __name__ == "__main__":
     test_load_names_from_file()
     test_class_label()
     test_hub_feature_flatten()
     test_feature_dict_str()
     test_feature_dict_repr()
+    test_classlabel_repr()
+    test_segmentation_repr()