added some minor changes

Author: naxatra2

src/deepforest/callbacks.py

@@ -9,6 +9,7 @@
 import glob
 import tempfile
 from deepforest.visualize import plot_results
+from deepforest.utilities import ColorPalette
 
 from pytorch_lightning import Callback
 from deepforest import dataset
@@ -59,6 +60,14 @@ def log_images(self, pl_module):
             selected_images = df.image_path.unique()[:self.n]
         df = df[df.image_path.isin(selected_images)]
 
+        # Ensure color is correctly assigned
+        if self.color is None:
+            num_classes = len(df["label"].unique())  # Determine number of classes
+            results_color = ColorPalette(
+                num_classes)  # Generate appropriate color palette
+        else:
+            results_color = self.color
+
         plot_results(results=df,
                      savedir=self.savedir,
                      results_color=self.color,