Display testing & fixes

Corrects color sampling in Display to avoid zero step and ensures image is always defined in display_frame. Adds comprehensive tests for Display, including headless environment setup, frame display, cutoff logic, window destruction, and color sampling safety.

Author: C-Achard

dlclive/display.py

@@ -7,7 +7,9 @@
 
 try:
     from tkinter import Label, Tk
+
     from PIL import ImageTk
+
     _TKINTER_AVAILABLE = True
 except ImportError:
     _TKINTER_AVAILABLE = False
@@ -33,9 +35,7 @@ class Display:
 
     def __init__(self, cmap="bmy", radius=3, pcutoff=0.5):
         if not _TKINTER_AVAILABLE:
-            raise ImportError(
-                "tkinter is not available. Display functionality requires tkinter. "
-            )
+            raise ImportError("tkinter is not available. Display functionality requires tkinter. ")
         self.cmap = cmap
         self.colors = None
         self.radius = radius
@@ -59,7 +59,9 @@ def set_display(self, im_size, bodyparts):
         self.lab.pack()
 
         all_colors = getattr(cc, self.cmap)
-        self.colors = all_colors[:: int(len(all_colors) / bodyparts)]
+        # Avoid 0 step
+        step = max(1, int(len(all_colors) / bodyparts))
+        self.colors = all_colors[::step]
 
     def display_frame(self, frame, pose=None):
         """
@@ -75,10 +77,10 @@ def display_frame(self, frame, pose=None):
         """
         if not _TKINTER_AVAILABLE:
             raise ImportError("tkinter is not available. Cannot display frames.")
-        
+
         im_size = (frame.shape[1], frame.shape[0])
+        img = Image.fromarray(frame)  # avoid undefined image if pose is None
         if pose is not None:
-            img = Image.fromarray(frame)
             draw = ImageDraw.Draw(img)
 
             if len(pose.shape) == 2:
@@ -91,33 +93,14 @@ def display_frame(self, frame, pose=None):
                 for j in range(pose.shape[1]):
                     if pose[i, j, 2] > self.pcutoff:
                         try:
-                            x0 = (
-                                pose[i, j, 0] - self.radius
-                                if pose[i, j, 0] - self.radius > 0
-                                else 0
-                            )
-                            x1 = (
-                                pose[i, j, 0] + self.radius
-                                if pose[i, j, 0] + self.radius < im_size[0]
-                                else im_size[1]
-                            )
-                            y0 = (
-                                pose[i, j, 1] - self.radius
-                                if pose[i, j, 1] - self.radius > 0
-                                else 0
-                            )
-                            y1 = (
-                                pose[i, j, 1] + self.radius
-                                if pose[i, j, 1] + self.radius < im_size[1]
-                                else im_size[0]
-                            )
+                            x0 = pose[i, j, 0] - self.radius if pose[i, j, 0] - self.radius > 0 else 0
+                            x1 = pose[i, j, 0] + self.radius if pose[i, j, 0] + self.radius < im_size[0] else im_size[1]
+                            y0 = pose[i, j, 1] - self.radius if pose[i, j, 1] - self.radius > 0 else 0
+                            y1 = pose[i, j, 1] + self.radius if pose[i, j, 1] + self.radius < im_size[1] else im_size[0]
                             coords = [x0, y0, x1, y1]
-                            draw.ellipse(
-                                coords, fill=self.colors[j], outline=self.colors[j]
-                            )
+                            draw.ellipse(coords, fill=self.colors[j], outline=self.colors[j])
                         except Exception as e:
                             print(e)
-
         img_tk = ImageTk.PhotoImage(image=img, master=self.window)
         self.lab.configure(image=img_tk)
         self.window.update()

tests/conftest.py

@@ -0,0 +1,22 @@
+import pytest
+
+
+@pytest.fixture
+def headless_display_env(monkeypatch):
+    # Import module under test
+    from test_display import FakeLabel, FakePhotoImage, FakeTk
+
+    import dlclive.display as display_mod
+
+    # Force tkinter availability and patch UI components
+    monkeypatch.setattr(display_mod, "_TKINTER_AVAILABLE", True, raising=False)
+    monkeypatch.setattr(display_mod, "Tk", FakeTk, raising=False)
+    monkeypatch.setattr(display_mod, "Label", FakeLabel, raising=False)
+
+    # Patch ImageTk.PhotoImage
+    class FakeImageTkModule:
+        PhotoImage = FakePhotoImage
+
+    monkeypatch.setattr(display_mod, "ImageTk", FakeImageTkModule, raising=False)
+
+    return display_mod

tests/test_display.py

@@ -0,0 +1,120 @@
+import numpy as np
+import pytest
+
+
+class FakeTk:
+    def __init__(self):
+        self.titles = []
+        self.updated = 0
+        self.destroyed = False
+
+    def title(self, text):
+        self.titles.append(text)
+
+    def update(self):
+        self.updated += 1
+
+    def destroy(self):
+        self.destroyed = True
+
+
+class FakeLabel:
+    def __init__(self, window):
+        self.window = window
+        self.packed = False
+        self.configured = {}
+
+    def pack(self):
+        self.packed = True
+
+    def configure(self, **kwargs):
+        self.configured.update(kwargs)
+
+
+class FakePhotoImage:
+    def __init__(self, image=None, master=None):
+        self.image = image
+        self.master = master
+
+
+def test_display_init_raises_when_tk_unavailable(monkeypatch):
+    import dlclive.display as display_mod
+
+    monkeypatch.setattr(display_mod, "_TKINTER_AVAILABLE", False, raising=False)
+
+    with pytest.raises(ImportError):
+        display_mod.Display()
+
+
+def test_display_frame_creates_window_and_updates(headless_display_env):
+    display_mod = headless_display_env
+    disp = display_mod.Display(radius=3, pcutoff=0.5)
+
+    frame = np.zeros((100, 120, 3), dtype=np.uint8)
+    pose = np.array([[[10, 10, 0.9], [50, 50, 0.2]]])  # 1 animal, 2 bodyparts
+
+    disp.display_frame(frame, pose)
+
+    assert disp.window is not None
+    assert disp.lab is not None
+    assert disp.lab.packed is True
+    assert disp.window.updated == 1
+    assert "image" in disp.lab.configured  # configured with PhotoImage
+
+
+def test_display_draws_only_points_above_cutoff(headless_display_env, monkeypatch):
+    display_mod = headless_display_env
+    disp = display_mod.Display(radius=3, pcutoff=0.5)
+
+    frame = np.zeros((100, 100, 3), dtype=np.uint8)
+    pose = np.array(
+        [
+            [
+                [10, 10, 0.9],  # draw
+                [20, 20, 0.49],  # don't draw
+                [30, 30, 0.5001],  # draw (>=)
+            ]
+        ],
+        dtype=float,
+    )
+
+    ellipses = []
+
+    class DrawRecorder:
+        def ellipse(self, coords, fill=None, outline=None):
+            ellipses.append((coords, fill, outline))
+
+    monkeypatch.setattr(display_mod.ImageDraw, "Draw", lambda img: DrawRecorder())
+
+    disp.display_frame(frame, pose)
+
+    assert len(ellipses) == 2
+
+
+def test_destroy_calls_window_destroy(headless_display_env):
+    display_mod = headless_display_env
+    disp = display_mod.Display()
+
+    frame = np.zeros((10, 10, 3), dtype=np.uint8)
+    pose = np.array([[[5, 5, 0.9]]])
+
+    disp.display_frame(frame, pose)
+    disp.destroy()
+
+    assert disp.window.destroyed is True
+
+
+def test_set_display_color_sampling_safe(headless_display_env, monkeypatch):
+    display_mod = headless_display_env
+
+    # Provide a fixed colormap list
+    class FakeCC:
+        bmy = [(1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), (0, 1, 1), (1, 0, 1)]
+
+    monkeypatch.setattr(display_mod, "cc", FakeCC)
+
+    disp = display_mod.Display(cmap="bmy")
+    disp.set_display(im_size=(100, 100), bodyparts=3)
+
+    assert disp.colors is not None
+    assert len(disp.colors) >= 3