Merge pull request #4 from TravisWheelerLab/develop

isaacrobinson2000 · web-flow · commit fb44c10b1bb5 · 2024-01-09T14:02:32.000-08:00
New Release: v0.0.9
diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md
@@ -1,6 +1,5 @@
-Changes for this version of diplomat:
- - Improved video player widget, with lag reduction when backtracking.
- - Two new UI labelers for source only placement of points.
- - Less lag for plots on the side in the UI when scrolling.
- - UI moved to it's own module and added to documentation.
- - Several additional minor bug fixes.
+Changes for this version of DIPLOMAT:
+ - Fixed installation process for DIPLOMAT with SLEAP on windows by adding keras dependency.
+ - Added support for changing some UI appearance settings to DIPLOMAT's supervised and tweak UI.
+ - Make SLEAP frontend error out unless a user explicitly passes a number of outputs parameter.
+ - Improved point rendering in the UI (proper alpha transparency support and smoother point rendering).
diff --git a/conda-environments/DIPLOMAT-SLEAP-CPU.yaml b/conda-environments/DIPLOMAT-SLEAP-CPU.yaml
@@ -37,6 +37,7 @@ dependencies:
   - conda-forge::seaborn
   - tensorflow >=2.6.0,<2.11  # No windows GPU support for >2.10
   - tensorflow-hub  # Pinned in meta.yml, but no problems here... yet
+  - keras
 
   # Packages required by tensorflow to find/use GPUs
   - conda-forge::cudatoolkit ==11.3.1
diff --git a/conda-environments/DIPLOMAT-SLEAP.yaml b/conda-environments/DIPLOMAT-SLEAP.yaml
@@ -37,6 +37,7 @@ dependencies:
   - conda-forge::seaborn
   - tensorflow-gpu >=2.6.0,<2.11  # No windows GPU support for >2.10
   - tensorflow-hub  # Pinned in meta.yml, but no problems here... yet
+  - keras
 
   # Packages required by tensorflow to find/use GPUs
   - conda-forge::cudatoolkit ==11.3.1
diff --git a/diplomat/__init__.py b/diplomat/__init__.py
@@ -2,7 +2,7 @@
 A tool providing multi-animal tracking capabilities on top of other Deep learning based tracking software.
 """
 
-__version__ = "0.0.8"
+__version__ = "0.0.9"
 # Can be used by functions to determine if diplomat was invoked through it's CLI interface.
 CLI_RUN = False
 
diff --git a/diplomat/frontends/sleap/predict_frames_sleap.py b/diplomat/frontends/sleap/predict_frames_sleap.py
@@ -65,7 +65,8 @@ def analyze_frames(
     """
     import sleap
     batch_size = _get_default_value(sleap.load_model, "batch_size", 4) if (batch_size is None) else batch_size
-    num_outputs = 1 if (num_outputs is None) else num_outputs
+    if(num_outputs is None):
+        raise ValueError("'num_outputs' is not set! Please set it to the number of bodies you are tracking.")
 
     print("Loading Config...")
     config = _load_config(_paths_to_str(config))[0]
diff --git a/diplomat/frontends/sleap/predict_videos_sleap.py b/diplomat/frontends/sleap/predict_videos_sleap.py
@@ -59,7 +59,8 @@ def analyze_videos(
 
     import sleap
     batch_size = _get_default_value(sleap.load_model, "batch_size", 4) if(batch_size is None) else batch_size
-    num_outputs = 1 if(num_outputs is None) else num_outputs
+    if(num_outputs is None):
+        raise ValueError("'num_outputs' is not set! Please set it to the number of bodies you are tracking.")
 
     print("Loading Model...")
     model = sleap.load_model(_paths_to_str(config), batch_size=batch_size)
diff --git a/diplomat/predictors/supervised_sfpe/supervised_segmented_frame_pass_engine.py b/diplomat/predictors/supervised_sfpe/supervised_segmented_frame_pass_engine.py
@@ -595,6 +595,14 @@ def _on_end(self, progress_bar: ProgressBar) -> Optional[Pose]:
             self._frame_holder.metadata["segments"] = self._segments.tolist()
             self._frame_holder.metadata["segment_scores"] = self._segment_scores.tolist()
         else:
+            progress_bar.reset(self._frame_holder.num_frames * self._frame_holder.num_bodyparts)
+            progress_bar.message("Restoring Partial Frames")
+            for frame_list in self._frame_holder.frames:
+                for frame in frame_list:
+                    if(frame.frame_probs is None):
+                        frame.frame_probs = frame.src_data.probs[:]
+                    progress_bar.update()
+
             self._width = self._frame_holder.metadata.width
             self._height = self._frame_holder.metadata.height
             self._resolve_frame_orderings(progress_bar)
diff --git a/diplomat/wx_gui/fpe_editor.py b/diplomat/wx_gui/fpe_editor.py
@@ -12,12 +12,14 @@
 from diplomat.utils.colormaps import iter_colormap
 from diplomat.utils.track_formats import to_diplomat_table, save_diplomat_table
 from diplomat.wx_gui.id_swap_dialog import IdSwapDialog
+from diplomat.wx_gui.labeler_lib import SettingCollection
 from diplomat.wx_gui.point_edit import PointEditor, PointViewNEdit, PoseLabeler
 from diplomat.wx_gui.progress_dialog import FBProgressDialog
 from diplomat.wx_gui.score_lib import ScoreEngine, ScoreEngineDisplayer
 from diplomat.wx_gui.scroll_image_list import ScrollImageList
 from diplomat.processing import Pose, ProgressBar
 from diplomat.wx_gui.helpdialog import HelpDialog
+from diplomat.wx_gui.settings_dialog import SettingsDialog
 from diplomat.wx_gui.video_player import VideoController
 from wx.lib.scrolledpanel import ScrolledPanel
 from collections import deque
@@ -435,6 +437,9 @@ def _build_toolbar(self):
         export_csv_bmp = icons.to_wx_bitmap(icons.SAVE_TRACKS_ICON, icons.SAVE_TRACKS_SIZE,
                                             self.GetForegroundColour(), self.TOOLBAR_ICON_SIZE)
 
+        settings_bmp = icons.to_wx_bitmap(icons.SETTINGS_ICON, icons.SETTINGS_SIZE,
+                                          self.GetForegroundColour(), self.TOOLBAR_ICON_SIZE)
+
         spin_ctrl = wx.SpinCtrl(self._toolbar, min=1, max=50, initial=PointViewNEdit.DEF_FAST_MODE_SPEED_FRACTION)
         spin_ctrl.SetMaxSize(wx.Size(-1, self.TOOLBAR_ICON_SIZE[1]))
         spin_ctrl.Bind(wx.EVT_SPINCTRL, self._on_spin)
@@ -474,8 +479,10 @@ def _build_toolbar(self):
         self._toolbar.AddTool(self._save)
         self._toolbar.AddSeparator()
 
-        self._turtle = self._toolbar.CreateTool(wx.ID_ANY, "Edit CTRL Speed: ", turtle_bmp, turtle_bmp,
-                                                shortHelp="Modify the labeling speed when CTRL Key is pressed (fast labeling mode).")
+        self._turtle = self._toolbar.CreateTool(
+            wx.ID_ANY, "Edit CTRL Speed: ", turtle_bmp, turtle_bmp,
+            shortHelp="Modify the labeling speed when CTRL Key is pressed (fast labeling mode)."
+        )
         self._toolbar.AddTool(self._turtle)
         self._toolbar.EnableTool(self._turtle.GetId(), False)
         self._toolbar.AddControl(spin_ctrl)
@@ -489,6 +496,10 @@ def _build_toolbar(self):
                                                        shortHelp="Export current tracks to a csv file from the UI.")
         self._toolbar.AddTool(self._export_to_csv)
 
+        self._visual_settings = self._toolbar.CreateTool(wx.ID_ANY, "Visual Settings", settings_bmp,
+                                                         shortHelp="Adjust some visual settings of the editor.")
+        self._toolbar.AddTool(self._visual_settings)
+
         self._help = self._toolbar.CreateTool(wx.ID_ANY, "Help", help_bmp, shortHelp="Display the help dialog.")
         self._toolbar.AddTool(self._help)
 
@@ -506,6 +517,7 @@ def _build_toolbar(self):
             self._turtle,
             self._export_btn,
             self._export_to_csv,
+            self._visual_settings,
             self._help
         ]
         self._bitmaps = [
@@ -519,6 +531,7 @@ def _build_toolbar(self):
             turtle_bmp,
             export_bmp,
             export_csv_bmp,
+            settings_bmp,
             help_bmp
         ]
 
@@ -697,9 +710,32 @@ def on_tool(self, evt: wx.CommandEvent):
             self._move_to_poor_label(False)
         elif(evt.GetId() == self._export_to_csv.GetId()):
             self._save_to_csv()
+        elif(evt.GetId() == self._visual_settings.GetId()):
+            self._change_visual_settings()
         elif((self._identity_swapper is not None) and (evt.GetId() == self._swap_id.GetId())):
             self._display_id_swap_dialog()
 
+    def _change_visual_settings(self):
+        from diplomat.wx_gui.labeler_lib import Slider, FloatSpin
+        from diplomat.wx_gui.settings_dialog import DropDown
+        from matplotlib import colormaps
+        point_video_viewer = self.video_player.video_viewer
+
+        sorted_colormaps = sorted(colormaps)
+
+        with SettingsDialog(self, title="Visual Settings", settings=SettingCollection(
+            colormap=DropDown([point_video_viewer.get_colormap()] + sorted_colormaps, ["CURRENT"] + sorted_colormaps),
+            point_radius=FloatSpin(1, 1000, point_video_viewer.get_point_radius(), increment=1, digits=0),
+            point_alpha=FloatSpin(0, 1, point_video_viewer.get_point_alpha(), increment=0.01, digits=2),
+            plot_threshold=FloatSpin(0, 1, point_video_viewer.get_plot_threshold(), increment=0.001, digits=3),
+            line_thickness=Slider(1, 10, point_video_viewer.get_line_thickness())
+        )) as dlg:
+            if(dlg.ShowModal() == wx.ID_OK):
+                for k, v in dlg.get_values().items():
+                    getattr(point_video_viewer, f"set_{k}")(v)
+                self.Refresh()
+                self.Update()
+
     def _save_to_csv(self):
         with wx.FileDialog(self, "Select FrameStore Save Location",
                            wildcard="CSV File (*.csv)",
diff --git a/diplomat/wx_gui/icons.py b/diplomat/wx_gui/icons.py
@@ -197,6 +197,31 @@ def _dump_pic(file):
 3NrxABagAUKi4lIyssyEFWMDrNy8IAMkOAgrxQ6YOXj5BXlZyNUOAuS2VgYSAAD00GPz
 """
 
+SETTINGS_SIZE = (64, 64)
+SETTINGS_ICON = b"""
+eJzll11oFFcUgG/WNSEBKUJqjGti0gi1IoXaSOJGLNo+WSTZTX8otS+aGHzoj8aArdDKsomhD6WF
+UgMVXWkL27R9K8UaLLTrg5jfTZHGNjQ2G92wuzFKfpowd+7pmTtzZ+/MzizbvHpeZu495zv37vm5
+c5eQx0CKvP8TqH7nG480PDAfaS4tGPa0/MoAXpFmoiqFf6NFBfK/AAWgw9mJCgVQ/ix0/TDTzKHF
+nDinTdCvC+Wf5jhN+ozxzmXu8FShPImr3EGitQQHpW9lKHf4QsH8ZzpAYfnOyOQKcG8AdQXzYYOw
+yrMF8z9TJ76zUNzPnHD1QaVpUeTJg9clHZcHdaJKLHDz/RzqyPWgXuR+FxzjmW7VDGq+B6Y8b8Mr
+5hkkQ43V/j7FgqdS8g5g8KOTUQWjSyfWW/koTjLGbYSkL7dUFRNSXBW4nM56AMP92xZ8d07IEkfX
+ZdXeYzM2NXuwQeZ/tPFqyNauZWGLhcrOyFrPl8xSMgu8d9btO98fi/X3NPFsBRelSK68aovfoVnJ
+/4IWXW+HuefEcS019QsmPpVbi0/0gfCgaqvXjZuxxEf8KZwKGGN6d3MOjvKecB/CQVMGAz14qt7n
+23N6CF8zfpwMGwb7nXBCYnpqZsoIqU0BvX9YKJqxpOa2YytPc4MpZ1xs4Cj+9nGgo9lyJ5VjFOIY
+gzaeumsu/GscT2PeO3D1Slm1BXfQjn4zmsWgC3+C85cwcRj5w1ZdC2YB0xjRLB65dN8nPL7NhOwD
+OmRXDquAIQzwJRqd+d95AqsJOe9wWJwG6Mbe4/n8whHfrxc2tsx3AHvs2gaAb7GZ+BLqAQd86xRP
+XwpfbwD4ctQAMXzwTqSroe029ebOh3r20ziIOfO/4SOj55jB/I0LAdFju84OM1GsDM/8fuf9Rwkp
+yXYJo1CvqxpB7j6MX49D/LoAwkb8TNmkq560TGLzNIGak78RFfZiE8uWq+KD/EiejeBxkOBlIAsm
+fhrL5opkyMaE8pp8ZKaxzo/j19NSv75ZCm1Yv3OSofq50H6gB88IwjG0iwMd2yLhcQqj2BftxsrK
+0MBQir0p1Dt6e0++/tIbX6ncwz3s3xrs36R5AQji6nP4+SxN6HzfRm22OOdS9PIqTw/Gmfi182O4
+q8Hna+gawdc0Bk/LiyYf2zlTOvXtBfC1Li6fX6O1ONXKvdOR9a689w+OLGiF4W1PiEhNt2nfgnrj
+AD7oihsHDMCitgPi8XdHY7Fo917e7606zsbz4KRc5DZcZlWU9RiVqzg3r5BlseeZNim83vZ7Yl65
+kA/fKH1FMpHgNmymkm3BK1LZsL/z3R1OgE0s329dOtzxTRnHm5NF1OUdbrjnqjvOFMX4bXS8xIX/
+0DCAiYE7AFIoGGO3L12Mi7tDrzO+Ia2fjpPa4bP1TNJ0QJdC1ZrBoTnNAU26/YCd2tWJ/lVuuPtB
+bHhanJbP4FVWvVvrGr+a20ydNdXFcWO/L5oGuxdhoNwVxw69+HBXdnREx+9LfxwOnsuXfpQK6d2n
+J+yn/IS7FK1ovBJZK08mefg+XTN/lWfg3TXzZ6eXlv659dya+cdA/gPgZyV7
+"""
+
 
 def to_wx_bitmap(icon_bytes: bytes, icon_size: Tuple[int, int], fg_color: wx.Colour, bitmap_size: Tuple[int, int] = None):
     """
@@ -227,7 +252,7 @@ def _main():
     # Tests icon loading code above by displaying the help icon in a wx Frame...,
     app = wx.App()
     frame = wx.Frame(None, title="Image Test!")
-    bitmap = to_wx_bitmap(SAVE_TRACKS_ICON, SAVE_TRACKS_SIZE, frame.GetForegroundColour())
+    bitmap = to_wx_bitmap(SETTINGS_ICON, SETTINGS_SIZE, frame.GetForegroundColour())
     icon = wx.StaticBitmap(frame, wx.ID_ANY, bitmap)
     frame.Show(1)
     app.MainLoop()
diff --git a/diplomat/wx_gui/labeler_lib.py b/diplomat/wx_gui/labeler_lib.py
@@ -325,7 +325,7 @@ def __init__(
         self,
         *args,
         title: str = "Advanced Settings",
-        collapsable = True,
+        collapsable: bool = True,
         **kwargs
     ):
         """
@@ -387,9 +387,11 @@ def _force_layout_fix(self, evt):
         # This is stupid, but by setting the sizer again we force the
         # window to resize to fit everything...
         self.SetSizerAndFit(self._main_layout)
+        self.Layout()
         # This is also stupid...
         window = wx.GetTopLevelParent(self)
         w, h = window.GetSize()
+        window.Layout()
         window.SetSize(w + 1, h + 1)
         window.SetSize(w, h)
 
diff --git a/diplomat/wx_gui/point_edit.py b/diplomat/wx_gui/point_edit.py
@@ -342,6 +342,9 @@ def on_draw(self, dc: wx.DC):
         # Call superclass draw to draw the video...
         super().on_draw(dc)
 
+        if(not isinstance(dc, wx.GCDC)):
+            dc = wx.GCDC(dc)
+
         width, height = self.GetClientSize()
         if((not width) or (not height)):
             return
@@ -366,7 +369,7 @@ def on_draw(self, dc: wx.DC):
             if(np.isnan(x) or np.isnan(y)):
                 continue
 
-            wx_color = wx.Colour(*color)
+            wx_color = wx.Colour(*color[:3], alpha=int(255 * self._point_alpha))
 
             if(prob < self._plot_threshold):
                 dc.SetBrush(wx.TRANSPARENT_BRUSH)
diff --git a/diplomat/wx_gui/settings_dialog.py b/diplomat/wx_gui/settings_dialog.py
@@ -0,0 +1,67 @@
+from typing import Any, Callable, List, Optional
+import wx
+from diplomat.processing import Config
+from diplomat.wx_gui.labeler_lib import SettingWidget, SettingCollection, SettingCollectionWidget
+
+
+class DropDown(SettingWidget):
+    def __init__(self, options: List[Any], option_names: Optional[List[str]] = None, default: int = 0, **kwargs):
+        if(len(options) == 0):
+            raise ValueError("No options offered!")
+        if(option_names is None):
+            option_names = [str(o) for o in options]
+        if(len(option_names) != len(options)):
+            raise ValueError("Options and name arrays don't have the same length.")
+        self._options = list(options)
+        self._option_names = option_names
+        if(not (0 <= default < len(options))):
+            raise ValueError("Default index is out of bounds!")
+        self._default = int(default)
+        self._kwargs = kwargs
+        self._hook = None
+        self._value = default
+
+    def set_hook(self, hook: Callable[[str], None]):
+        self._hook = hook
+
+    def get_new_widget(self, parent=None) -> wx.Control:
+        text_list = wx.Choice(parent, choices=self._option_names, style=wx.LB_SINGLE, **self._kwargs)
+        text_list.SetSelection(self._default)
+
+        def val_change(evt):
+            sel = text_list.GetSelection()
+            if(sel == wx.NOT_FOUND):
+                text_list.SetSelection(self._default)
+                self._value = self._default
+            else:
+                self._value = sel
+
+            if(self._hook is not None):
+                self._hook(self._options[self._value])
+
+        text_list.Bind(wx.EVT_CHOICE, val_change)
+
+        return text_list
+
+    def get_value(self) -> Any:
+        return self._options[self._value]
+
+
+class SettingsDialog(wx.Dialog):
+    def __init__(self, *args, title: str = "Settings", settings: SettingCollection, **kwargs):
+        super().__init__(*args, title=title, style=wx.DEFAULT_DIALOG_STYLE | wx.RESIZE_BORDER, **kwargs)
+
+        self._parent_layout = wx.BoxSizer(wx.VERTICAL)
+
+        self._settings = settings
+        self._setting_widget = SettingCollectionWidget(self, title=title, collapsable=False)
+        self._setting_widget.set_setting_collection(settings)
+        self._parent_layout.Add(self._setting_widget, proportion=1, flag=wx.EXPAND | wx.ALL)
+
+        self._buttons = self.CreateButtonSizer(wx.OK | wx.CANCEL)
+        self._parent_layout.Add(self._buttons, proportion=0, flag=wx.EXPAND | wx.ALL)
+
+        self.SetSizerAndFit(self._parent_layout)
+
+    def get_values(self) -> Config:
+        return self._settings.get_values()
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -6,6 +6,7 @@
 sys.path.insert(0, str(Path(__file__).resolve().parent / "ext"))
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
 
+
 def _get_version() -> str:
     with mock(autodoc_mock_imports):
         import diplomat
