Allow user to set grayscale when replacing videos (mp4/avi only)

sleap/gui/commands.py

@@ -35,7 +35,7 @@ class which inherits from `AppCommand` (or a more specialized class such as
 
 from enum import Enum
 from glob import glob
-from pathlib import PurePath
+from pathlib import PurePath, Path
 from typing import Callable, Dict, Iterator, List, Optional, Type, Tuple
 
 import numpy as np
@@ -51,6 +51,7 @@ class which inherits from `AppCommand` (or a more specialized class such as
 from sleap.io.dataset import Labels
 from sleap.gui.dialogs.delete import DeleteDialog
 from sleap.gui.dialogs.importvideos import ImportVideos
+from sleap.gui.dialogs.query import QueryDialog
 from sleap.gui.dialogs.filedialog import FileDialog
 from sleap.gui.dialogs.missingfiles import MissingFilesDialog
 from sleap.gui.dialogs.merge import MergeDialog
@@ -1488,30 +1489,106 @@ def do_action(context: CommandContext, params: dict):
 
 
 class ReplaceVideo(EditCommand):
-    topics = [UpdateTopic.video]
-
-    @staticmethod
-    def do_action(context: CommandContext, params: dict):
-        new_paths = params["new_video_paths"]
-
-        for video, new_path in zip(context.labels.videos, new_paths):
-            if new_path != video.backend.filename:
-                video.backend.filename = new_path
-                video.backend.reset()
+    topics = [UpdateTopic.video, UpdateTopic.frame]
 
     @staticmethod
-    def ask(context: CommandContext, params: dict) -> bool:
+    def ask_and_do(context: CommandContext, params: dict) -> bool:
         """Shows gui for replacing videos in project."""
-        paths = [video.backend.filename for video in context.labels.videos]
 
-        okay = MissingFilesDialog(filenames=paths, replace=True).exec_()
+        # Warn user: newly added labels will be discarded if project is not saved
+        if not context.state["filename"] or context.state["has_changes"]:
+            QMessageBox(
+                text=("You have unsaved changes. Please save before replacing videos.")
+            ).exec_()
+            return False
 
+        # Select the videos we want to swap
+        old_paths = [video.backend.filename for video in context.labels.videos]
+        paths = list(old_paths)
+        okay = MissingFilesDialog(filenames=paths, replace=True).exec_()
         if not okay:
             return False
 
-        params["new_video_paths"] = paths
+        # Only return an import list for videos we swap
+        new_paths = [
+            (path, video_idx)
+            for video_idx, (path, old_path) in enumerate(zip(paths, old_paths))
+            if path != old_path
+        ]
+
+        new_paths = []
+        old_videos = []
+        all_videos = context.labels.videos
+        for video_idx, (path, old_path) in enumerate(zip(paths, old_paths)):
+            if path != old_path:
+                new_paths.append(path)
+                old_videos.append(all_videos[video_idx])
 
-        return True
+        params["import_list"] = zip(ImportVideos().ask(filenames=new_paths), old_videos)
+
+        # This is the beginning of the "do" part of the funtion, but has a GUI pop-up.
+
+        def _trim_labeled_frames(lfs_to_remove):
+            """Trim labels past new video length."""
+            for lf in lfs_to_remove:
+                context.labels.remove_frame(lf)
+
+        def _reset_video_backend(video, filename, grayscale):
+            """Reset video back to original if operation aborted."""
+            video.backend.reset(
+                filename=filename,
+                grayscale=grayscale,
+            )
+
+        import_list = params["import_list"]
+
+        for import_item, video in import_list:
+            import_params = import_item["params"]
+            old_path = video.backend.filename
+            old_grayscale = video.backend.grayscale
+            new_path = import_params["filename"]
+
+            # TODO: Will need to create a new backend if import has different extension.
+            # Currently reset is only functional for MediaVideo backend.
+            # FIXME: Ideally we would only mess with the backend after getting approval
+            # ...but, use this to read in the number of frames. Could use cv2 here....
+            if new_path != video.backend.filename:
+                _reset_video_backend(video, new_path, import_params)
+
+            # Remove frames in video past last frame index
+            last_vid_frame = video.last_frame_idx
+            lfs: List[LabeledFrame] = list(context.labels.get(video))
+            if lfs is not None:
+                lfs.sort(key=lambda lf: lf.frame_idx)
+                last_lf_frame = lfs[-1].frame_idx
+                lfs = [lf for lf in lfs if lf.frame_idx > last_vid_frame]
+
+                # Warn users that labels will be removed if proceed
+                if last_lf_frame > last_vid_frame:
+                    message = (
+                        "<p><strong>Warning:</strong> The replacement video is shorter "
+                        f"than the frame index of {len(lfs)} labeled frames.</p>"
+                        f"<p><em>Current video</em>: <b>{Path(old_path).name}</b> "
+                        f"(last label at frame {last_lf_frame})<br>"
+                        f"<em>Replacement video</em>: <b>{Path(video.filename).name}"
+                        f"</b> ({last_vid_frame} frames)</p>"
+                        f"<p>Replace video (and remove {len(lfs)} labeled frames past "
+                        f"frame {last_vid_frame})?</p>"
+                    )
+                    query = QueryDialog(
+                        "Replace Video: Truncate Labeled Frames", message, context.app
+                    )
+                    query.accepted.connect(lambda: _trim_labeled_frames(lfs))
+                    # FIXME: Ideally we do not mess with backend before approval
+                    query.rejected.connect(
+                        lambda: _reset_video_backend(video, old_path, old_grayscale)
+                    )
+                    query.exec_()
+                else:
+                    _trim_labeled_frames(lfs)
+
+            # Update seekbar and video length through callbacks
+            context.state.emit("video")
 
 
 class RemoveVideo(EditCommand):

sleap/gui/dialogs/query.py

@@ -0,0 +1,31 @@
+"""
+Generic module to ask user permission to complete an action.
+"""
+
+from PySide2.QtWidgets import (
+    QDialog,
+    QDialogButtonBox,
+    QVBoxLayout,
+    QLabel,
+)
+
+
+class QueryDialog(QDialog):
+    def __init__(self, title: str, message: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.user_response = False
+
+        self.setWindowTitle(title)
+
+        QBtn = QDialogButtonBox.Ok | QDialogButtonBox.Cancel
+
+        self.buttonBox = QDialogButtonBox(QBtn)
+        self.buttonBox.accepted.connect(self.accept)
+        self.buttonBox.rejected.connect(self.reject)
+
+        self.layout = QVBoxLayout()
+        message = QLabel(message)
+        self.layout.addWidget(message)
+        self.layout.addWidget(self.buttonBox)
+        self.setLayout(self.layout)

sleap/io/video.py

@@ -461,9 +461,23 @@ def dtype(self):
         """See :class:`Video`."""
         return self.test_frame.dtype
 
-    def reset(self):
+    def reset(self, filename: str = None, grayscale: bool = None, bgr: bool = None):
         """Reloads the video."""
-        self._reader_ = None
+        if filename is not None:
+            self.filename = filename
+            self._test_frame_ = None  # Test frame does not depend on num channels
+
+        if grayscale is not None:
+            self.grayscale = grayscale
+            self._detect_grayscale = False
+        else:
+            self._detect_grayscale = True
+
+        if bgr is not None:
+            self.bgr = bgr
+
+        if (filename is not None) or (grayscale is not None):
+            self._reader_ = None  # Reader depends on both filename and grayscale
 
     def get_frame(self, idx: int, grayscale: bool = None) -> np.ndarray:
         """See :class:`Video`."""

tests/io/test_video.py

@@ -1,3 +1,4 @@
+from cgitb import reset
 import pytest
 import os
 import h5py
@@ -405,3 +406,50 @@ def test_load_video():
     video = load_video(TEST_SMALL_CENTERED_PAIR_VID)
     assert video.shape == (1100, 384, 384, 1)
     assert video[:3].shape == (3, 384, 384, 1)
+
+
+def test_reset_video(small_robot_mp4_vid: Video):
+    def assert_video_params(
+        filename: str = None,
+        grayscale: bool = None,
+        bgr: bool = small_robot_mp4_vid.backend.bgr,
+        reset: bool = False,
+    ):
+        assert video.backend.filename == filename
+        assert video.backend.grayscale == grayscale
+        assert video.backend.bgr == bgr
+        if reset:
+            assert video.backend._reader_ == None
+            assert video.backend._test_frame_ == None
+            assert video.backend._detect_grayscale == bool(grayscale is None)
+        # Getting the channels will assert some of the above are not None
+        assert video.backend.channels == 3 ** (not grayscale)
+
+    video = small_robot_mp4_vid
+    filename = video.backend.filename
+
+    # Get a frame to set the video parameters in the backend
+    video.get_frame(idx=0)
+    assert_video_params(filename=filename, grayscale=video.backend.grayscale)
+
+    # Test reset works for color to grayscale
+
+    # Reset the backend
+    video.backend.reset(filename=filename, grayscale=True)
+    assert_video_params(filename=filename, grayscale=True, reset=True)
+
+    # Get a frame to test that reset parameters persist (namely grayscale and channels)
+    frame = video.get_frame(idx=0)
+    assert frame.shape[2] == 1
+    assert_video_params(filename=filename, grayscale=True)
+
+    # Test reset works for grayscale to color
+
+    # Reset the backend
+    video.backend.reset(filename=filename, grayscale=False)
+    assert_video_params(filename=filename, grayscale=False, bgr=True, reset=True)
+
+    # Get a frame to test that reset parameters persist (namely grayscale and channels)
+    frame = video.get_frame(idx=0)
+    assert frame.shape[2] == 3
+    assert_video_params(filename=filename, grayscale=False)