Merge pull request #3037 from activeloopai/fix_apllete

Fix palette
activeloopai · Feb 24, 2025 · 0af67e2 · 0af67e2
2 parents 1713dc2 + 3f43dec
commit 0af67e2
Showing 1 changed file with 16 additions and 1 deletion.
diff --git a/deeplake/integrations/mmlab/segmentation/mmseg_.py b/deeplake/integrations/mmlab/segmentation/mmseg_.py
@@ -50,6 +50,17 @@
 )
 
 
+def __generate_palette(num_classes):
+    """Function to generate a random but distinguishable color palette"""
+    import random
+
+    random.seed(42)
+    return [
+        [random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)]
+        for _ in range(num_classes)
+    ]
+
+
 def build_dataloader(
     dataloader: Union[DeepLakeDataLoader, Dict],
     seed: Optional[int] = None,
@@ -200,7 +211,11 @@ def build_dataloader(
             persistent_workers=persistent_workers,
         )
     )
-    loader.dataset.__setattr__("metainfo", {"classes": classes})
+
+    loader.dataset.__setattr__(
+        "metainfo",
+        {"classes": tuple(classes), "palette": __generate_palette(len(classes))},
+    )
 
     if init_fn:
         loader.worker_init_fn = init_fn